diff --git a/documentation/source/files/technical_description.rst b/documentation/source/files/technical_description.rst index 75ffe6e..7270c02 100644 --- a/documentation/source/files/technical_description.rst +++ b/documentation/source/files/technical_description.rst @@ -94,6 +94,20 @@ The computational flow is as follows: Sketch of the computational flow. + +Generating Symbolic Equations +----------------------------- +Using `Sympy`_ qgs offers the functionality to return the ordinary differential equations of the projected model as a string, with any parameters the user chooses to be returned as varibales in the equations. In addition, the resulting equations can be returned already formatted in the programming language of the users' choice. This allows the qgs framework to feed directly into pipelines in other programming languages. Currently the framework can return the model equations in the following languages: + +* Python +* Julia +* Fortran-90 +* Mathematica +* AUTO-p07 continuation software + +This also allows the user to specify their own integration method for solving the model equations in python. + + Additional technical information -------------------------------- @@ -103,6 +117,8 @@ Additional technical information * qgs has a `tangent linear model`_ optimized to run ensembles of initial conditions as well, with a broadcast integration of the tangent model thanks to `Numpy`_. +* The symbolic output functionality of the qgs model relies on `Sympy`_ to perform the tensor calculations. This library is significantly slower than the numerical equivilent and as a result it is currently only feasible to generate the symbolic model equations for model resolutions up to :math:`4x4`. + References ---------- diff --git a/documentation/source/files/user_guide.rst b/documentation/source/files/user_guide.rst index 390d57a..f73e18a 100644 --- a/documentation/source/files/user_guide.rst +++ b/documentation/source/files/user_guide.rst @@ -462,7 +462,33 @@ This concludes the initialization of qgs, the function :meth:`f` hence produced An example of the construction exposed here, along with plots of the trajectories generated, can be found in the section :ref:`files/examples/manual:Manual setting of the basis`. See also the following section for the possible usages. -4. Using qgs (once configured) + +4. Symbolic outputs +------------------- + +If the user wants to generate the model tendencies with non-fixed parameters, use the tendencies in another programming language, or use their own integrator in python, the qgs framework can use `Sympy`_ to proform the above calculations symbolically rather than numerically. + +To return the tendencies including specified parameter values, or to format the tendencies in another programming language, the qgs model is configured as shown in section :ref:`files/user_guide:Configuration of qgs`, however the **Symbolic** inner product is always used in this pipeline. To create the symbolic tendencies, the instance of :class:`.QgParams` is passed to the function :func:`.create_symbolic_tendencies`: + +.. code:: ipython3 + + from qgs.functions.symbolic_tendencies import create_symbolic_equations + + parameters = [model_parameters.par1, model_parameters.par2, model_parameters.par3] + + f = create_symbolic_equations(model_parameters, continuation_variables=parameters, language='python') + +The varibale :meth:`f` is a string of the model tendencies, formatted in the programming language specified by the keyword :meth:`language`. The model tendencies will contain the specified :meth:`continuation_variables` as free parameters. + +Currently the framework can format the equations in the following programming languages: +* :meth:`python` +* :meth:`julia` +* :meth:`fortran` +* :meth:`mathematica` +* :meth:`auto` + + +1. Using qgs (once configured) --------------------------------- Once the function :math:`\boldsymbol{f}` giving the model's tendencies has been obtained, it is possible to use it with @@ -482,7 +508,7 @@ the qgs built-in integrator to obtain the model's trajectories: Note that it is also possible to use other ordinary differential equations integrators available on the market, see for instance the :ref:`files/examples/diffeq_example:Example of DiffEqPy usage`. -4.1 Analyzing the model's output with diagnostics +5.1 Analyzing the model's output with diagnostics ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ .. include:: diagnostic.rst @@ -524,7 +550,7 @@ Note that it is also possible to use other ordinary differential equations integ More diagnostics will be implemented soon. -4.2 Toolbox +5.2 Toolbox ^^^^^^^^^^^ The toolbox regroups submodules to make specific analysis with the model and are available in the :mod:`qgs.toolbox` module. @@ -536,7 +562,7 @@ Presently, the list of submodules available is the following: More submodules will be implemented soon. -4.2.1 Lyapunov toolbox +5.2.1 Lyapunov toolbox """""""""""""""""""""" This module allows you to integrate the model and simultaneously obtain the *local* `Lyapunov vectors`_ and `exponents`_. @@ -555,10 +581,10 @@ See also :cite:`user-KP2012` for a description of these methods. Some example notebooks on how to use this module are available in the `notebooks/lyapunov <../../../../notebooks/lyapunov>`_ folder. -5. Developers information +6. Developers information ------------------------- -5.1 Running the test +6.1 Running the test ^^^^^^^^^^^^^^^^^^^^^ The model core tensors can be tested by running `pytest`_ in the main folder: :: diff --git a/model_test/test_aotensor_sym.py b/model_test/test_aotensor_sym.py new file mode 100644 index 0000000..d92c446 --- /dev/null +++ b/model_test/test_aotensor_sym.py @@ -0,0 +1,103 @@ + +import sys +import os + +path = os.path.abspath('./') +base = os.path.basename(path) +if base == 'model_test': + sys.path.extend([os.path.abspath('../')]) +else: + sys.path.extend([path]) + + +import unittest +import numpy as np + +from qgs.params.params import QgParams +from qgs.inner_products import symbolic +from qgs.tensors.qgtensor import QgsTensor +from qgs.tensors.symbolic_qgtensor import SymbolicQgsTensor + +from model_test.test_base_symbolic import TestBaseSymbolic + +real_eps = np.finfo(np.float64).eps + + +class TestSymbolicAOTensor(TestBaseSymbolic): + ''' + Test class for the Linear Symbolic Tensor + The tensor is tested against the reference file, and then the numerical tensor calculated in qgs. + ''' + + filename = 'test_aotensor.ref' + + def test_sym_against_ref(self): + self.check_lists_flt() + + def test_sym_against_num(self): + self.check_numerical_lists_flt() + + def symbolic_outputs(self, output_func=None): + + if output_func is None: + self.symbolic_values.clear() + tfunc = self.save_ip_symbolic + else: + tfunc = output_func + + params = QgParams({'rr': 287.e0, 'sb': 5.6e-8}) + params.set_atmospheric_channel_fourier_modes(2, 2, mode="symbolic") + params.set_oceanic_basin_fourier_modes(2, 4, mode="symbolic") + + # Setting MAOOAM default parameters + params.set_params({'kd': 0.04, 'kdp': 0.04, 'n': 1.5}) + + aip = symbolic.AtmosphericSymbolicInnerProducts(params, return_symbolic=True, make_substitution=True) + oip = symbolic.OceanicSymbolicInnerProducts(params, return_symbolic=True, make_substitution=True) + + aip.connect_to_ocean(oip) + + sym_aotensor = SymbolicQgsTensor(params=params, atmospheric_inner_products=aip, oceanic_inner_products=oip) + + subbed_tensor = sym_aotensor.sub_tensor() + + for coo, val in zip(subbed_tensor.keys(), subbed_tensor.values()): + _ip_string_format(tfunc, 'sym_aotensor', coo, val) + + def numerical_outputs(self, output_func=None): + + if output_func is None: + self.numerical_values.clear() + tfunc = self.save_ip_numeric + else: + tfunc = output_func + + params = QgParams({'rr': 287.e0, 'sb': 5.6e-8}) + params.set_atmospheric_channel_fourier_modes(2, 2, mode="symbolic") + params.set_oceanic_basin_fourier_modes(2, 4, mode="symbolic") + + # Setting MAOOAM default parameters + params.set_params({'kd': 0.04, 'kdp': 0.04, 'n': 1.5}) + + aip = symbolic.AtmosphericSymbolicInnerProducts(params, return_symbolic=False) + oip = symbolic.OceanicSymbolicInnerProducts(params, return_symbolic=False) + + aip.connect_to_ocean(oip) + + num_aotensor = QgsTensor(params=params, atmospheric_inner_products=aip, oceanic_inner_products=oip) + + for coo, val in zip(num_aotensor.tensor.coords.T, num_aotensor.tensor.data): + _ip_string_format(tfunc, 'num_aotensor', coo, val) + + +def _ip_string_format(func, symbol, indices, value): + if abs(value) >= real_eps: + s = symbol + for i in indices: + s += "["+str(i)+"]" + s += " = % .5E" % value + func(s) + + +if __name__ == "__main__": + unittest.main() diff --git a/model_test/test_aotensor_sym_dynT.py b/model_test/test_aotensor_sym_dynT.py new file mode 100644 index 0000000..852b2a6 --- /dev/null +++ b/model_test/test_aotensor_sym_dynT.py @@ -0,0 +1,97 @@ + +import sys +import os + +path = os.path.abspath('./') +base = os.path.basename(path) +if base == 'model_test': + sys.path.extend([os.path.abspath('../')]) +else: + sys.path.extend([path]) + + +import unittest +import numpy as np + +from qgs.params.params import QgParams +from qgs.inner_products import symbolic +from qgs.tensors.qgtensor import QgsTensorDynamicT +from qgs.tensors.symbolic_qgtensor import SymbolicQgsTensorDynamicT + +from model_test.test_base_symbolic import TestBaseSymbolic + +real_eps = np.finfo(np.float64).eps + +class TestSymbolicAOTensorDynT(TestBaseSymbolic): + ''' + Test class for the Dynamic T Symbolic Tensor + The tensor is tested against the reference file, and then the numerical tensor calculated in qgs. + ''' + + def test_sym_against_num(self): + self.check_numerical_lists_flt() + + def symbolic_outputs(self, output_func=None): + + if output_func is None: + self.symbolic_values.clear() + tfunc = self.save_ip_symbolic + else: + tfunc = output_func + + params = QgParams({'rr': 287.e0, 'sb': 5.6e-8}, dynamic_T=True) + params.set_atmospheric_channel_fourier_modes(2, 2, mode="symbolic") + params.set_oceanic_basin_fourier_modes(2, 4, mode="symbolic") + + # Setting MAOOAM default parameters + params.set_params({'kd': 0.04, 'kdp': 0.04, 'n': 1.5}) + + aip = symbolic.AtmosphericSymbolicInnerProducts(params, return_symbolic=True, make_substitution=True) + oip = symbolic.OceanicSymbolicInnerProducts(params, return_symbolic=True, make_substitution=True) + + if not aip.connected_to_ocean: + aip.connect_to_ocean(oip) + + sym_aotensor = SymbolicQgsTensorDynamicT(params=params, atmospheric_inner_products=aip, oceanic_inner_products=oip) + + subbed_tensor = sym_aotensor.sub_tensor() + + for coo, val in zip(subbed_tensor.keys(), subbed_tensor.values()): + _ip_string_format(tfunc, 'sym_aotensor', coo, val) + + def numerical_outputs(self, output_func=None): + + if output_func is None: + self.numerical_values.clear() + tfunc = self.save_ip_numeric + else: + tfunc = output_func + + params = QgParams({'rr': 287.e0, 'sb': 5.6e-8}, dynamic_T=True) + params.set_atmospheric_channel_fourier_modes(2, 2, mode="symbolic") + params.set_oceanic_basin_fourier_modes(2, 4, mode="symbolic") + + # Setting MAOOAM default parameters + params.set_params({'kd': 0.04, 'kdp': 0.04, 'n': 1.5}) + + aip = symbolic.AtmosphericSymbolicInnerProducts(params, return_symbolic=False) + oip = symbolic.OceanicSymbolicInnerProducts(params, return_symbolic=False) + + if not aip.connected_to_ocean: + aip.connect_to_ocean(oip) + + num_aotensor = QgsTensorDynamicT(params=params, atmospheric_inner_products=aip, oceanic_inner_products=oip) + + for coo, val in zip(num_aotensor.tensor.coords.T, num_aotensor.tensor.data): + _ip_string_format(tfunc, 'num_aotensor', coo, val) + +def _ip_string_format(func, symbol, indices, value): + if abs(value) >= real_eps: + s = symbol + for i in indices: + s += "["+str(i)+"]" + s += " = % .5E" % value + func(s) + +if __name__ == "__main__": + unittest.main() diff --git a/model_test/test_aotensor_sym_ground.py b/model_test/test_aotensor_sym_ground.py new file mode 100644 index 0000000..2ba2cbe --- /dev/null +++ b/model_test/test_aotensor_sym_ground.py @@ -0,0 +1,113 @@ + +import sys +import os + +path = os.path.abspath('./') +base = os.path.basename(path) +if base == 'model_test': + sys.path.extend([os.path.abspath('../')]) +else: + sys.path.extend([path]) + + +import unittest +import numpy as np + +from qgs.params.params import QgParams +from qgs.inner_products import symbolic +from qgs.tensors.qgtensor import QgsTensor +from qgs.tensors.symbolic_qgtensor import SymbolicQgsTensor + +from model_test.test_base_symbolic import TestBaseSymbolic + +real_eps = np.finfo(np.float64).eps + + +class TestSymbolicGroundTensor(TestBaseSymbolic): + ''' + Test class for the Linear Symbolic Tensor + The tensor is tested against the reference file, and then the numerical tensor calculated in qgs. + ''' + + filename = 'test_aotensor.ref' + + # def test_sym_against_ref(self): + # self.check_lists_flt() + + def test_sym_against_num(self): + self.check_numerical_lists_flt() + + def symbolic_outputs(self, output_func=None): + + if output_func is None: + self.symbolic_values.clear() + tfunc = self.save_ip_symbolic + else: + tfunc = output_func + + params = QgParams({'rr': 287.e0, 'sb': 5.6e-8}) + params.set_atmospheric_channel_fourier_modes(2, 2, mode="symbolic") + params.set_ground_channel_fourier_modes(2, 2, mode="symbolic") + + params.ground_params.set_orography(0.2, 1) + params.gotemperature_params.set_params({'gamma': 1.6e7, 'T0': 300}) + params.atemperature_params.set_params({'hlambda': 10, 'T0': 290}) + params.atmospheric_params.set_params({'sigma': 0.2, 'kd': 0.085, 'kdp': 0.02}) + C_g = 300 + params.atemperature_params.set_insolation(0.4 * C_g, 0) + params.gotemperature_params.set_insolation(C_g, 0) + + aip = symbolic.AtmosphericSymbolicInnerProducts(params, return_symbolic=True, make_substitution=True) + gip = symbolic.GroundSymbolicInnerProducts(params, return_symbolic=True, make_substitution=True) + + aip.connect_to_ground(gip, None) + + sym_aotensor = SymbolicQgsTensor(params=params, atmospheric_inner_products=aip, ground_inner_products=gip) + + subbed_tensor = sym_aotensor.sub_tensor() + + for coo, val in zip(subbed_tensor.keys(), subbed_tensor.values()): + _ip_string_format(tfunc, 'sym_aotensor', coo, val) + + def numerical_outputs(self, output_func=None): + + if output_func is None: + self.numerical_values.clear() + tfunc = self.save_ip_numeric + else: + tfunc = output_func + + params = QgParams({'rr': 287.e0, 'sb': 5.6e-8}) + params.set_atmospheric_channel_fourier_modes(2, 2, mode="symbolic") + params.set_ground_channel_fourier_modes(2, 2, mode="symbolic") + + params.ground_params.set_orography(0.2, 1) + params.gotemperature_params.set_params({'gamma': 1.6e7, 'T0': 300}) + params.atemperature_params.set_params({'hlambda': 10, 'T0': 290}) + params.atmospheric_params.set_params({'sigma': 0.2, 'kd': 0.085, 'kdp': 0.02}) + C_g = 300 + params.atemperature_params.set_insolation(0.4 * C_g, 0) + params.gotemperature_params.set_insolation(C_g, 0) + + aip = symbolic.AtmosphericSymbolicInnerProducts(params, return_symbolic=False) + gip = symbolic.GroundSymbolicInnerProducts(params, return_symbolic=False) + + aip.connect_to_ground(gip, None) + + num_aotensor = QgsTensor(params=params, atmospheric_inner_products=aip, ground_inner_products=gip) + + for coo, val in zip(num_aotensor.tensor.coords.T, num_aotensor.tensor.data): + _ip_string_format(tfunc, 'num_aotensor', coo, val) + + +def _ip_string_format(func, symbol, indices, value): + if abs(value) >= real_eps: + s = symbol + for i in indices: + s += "["+str(i)+"]" + s += " = % .5E" % value + func(s) + + +if __name__ == "__main__": + unittest.main() diff --git a/model_test/test_base_symbolic.py b/model_test/test_base_symbolic.py new file mode 100644 index 0000000..b540410 --- /dev/null +++ b/model_test/test_base_symbolic.py @@ -0,0 +1,92 @@ + +import os + +path = os.path.abspath('./') +base = os.path.basename(path) +if base == 'model_test': + fold = "" +else: + fold = 'model_test/' + +import unittest +import numpy as np + +real_eps = np.finfo(np.float64).eps + +class TestBaseSymbolic(unittest.TestCase): + + reference = list() + symbolic_values = list() + numerical_values = list() + folder = fold + + def load_ref_from_file(self): + self.reference.clear() + f = open(self.folder+self.filename, 'r') + buf = f.readlines() + + for l in buf: + self.reference.append(l[:-1]) + + f.close() + + def save_ip_symbolic(self, s): + self.symbolic_values.append(s) + + def save_ip_numeric(self, s): + self.numerical_values.append(s) + + def check_lists_flt(self): + self.symbolic_outputs() + self.load_ref_from_file() + for v, r in zip(list(reversed(sorted(self.symbolic_values))), list(reversed(sorted(self.reference)))): + self.assertTrue(self.match_flt(v, r), msg=v+' != '+r+' !!!') + + def check_lists(self, cmax=1): + self.symbolic_outputs() + self.load_ref_from_file() + for v, r in zip(list(reversed(sorted(self.symbolic_values))), list(reversed(sorted(self.reference)))): + self.assertTrue(self.match_str(v, r, cmax), msg=v+' != '+r+' !!!') + + def check_numerical_lists_flt(self): + self.symbolic_outputs() + self.numerical_outputs() + for v, r in zip(list(reversed(sorted(self.symbolic_values))), list(reversed(sorted(self.numerical_values)))): + self.assertTrue(self.match_flt(v, r), msg=v+' != '+r+' !!!') + + def check_numerical_lists(self, cmax=1): + self.symbolic_outputs() + self.numerical_outputs() + for v, r, in zip(list(reversed(sorted(self.symbolic_values))), list(reversed(sorted(self.numerical_values)))): + self.assertTrue(self.match_str(v, r, cmax), msg=v+' != '+r+' !!!') + + def symbolic_outputs(self): + pass + + def numerical_outputs(self): + pass + + @staticmethod + def match_flt(s1, s2, eps=real_eps): + + s1p = s1.split('=') + s2p = s2.split('=') + + v1 = float(s1p[1]) + v2 = float(s2p[1]) + + return abs(v1 - v2) < eps + + @staticmethod + def match_str(s1, s2, cmax=1): + + c = 0 + + for c1, c2 in zip(s1, s2): + if c1 != c2: + c += 1 + + if c > cmax: + return False + + return True diff --git a/notebooks/ground_heat.ipynb b/notebooks/ground_heat.ipynb index 168b078..f5b74db 100644 --- a/notebooks/ground_heat.ipynb +++ b/notebooks/ground_heat.ipynb @@ -44,7 +44,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "metadata": {}, "outputs": [], "source": [ @@ -53,7 +53,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "metadata": {}, "outputs": [], "source": [ @@ -62,7 +62,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 3, "metadata": {}, "outputs": [], "source": [ @@ -80,7 +80,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 4, "metadata": {}, "outputs": [], "source": [ @@ -96,7 +96,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 5, "metadata": {}, "outputs": [], "source": [ @@ -122,7 +122,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 6, "metadata": {}, "outputs": [], "source": [ @@ -144,7 +144,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 7, "metadata": {}, "outputs": [], "source": [ @@ -164,7 +164,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 8, "metadata": {}, "outputs": [], "source": [ @@ -174,9 +174,85 @@ }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Qgs v0.2.8 parameters summary\n", + "=============================\n", + "\n", + "General Parameters:\n", + "'dynamic_T': False,\n", + "'T4': False,\n", + "'time_unit': days,\n", + "'rr': 287.058 [J][kg^-1][K^-1] (gas constant of dry air),\n", + "'sb': 5.67e-08 [J][m^-2][s^-1][K^-4] (Stefan-Boltzmann constant),\n", + "\n", + "Scale Parameters:\n", + "'scale': 5000000.0 [m] (characteristic space scale (L*pi)),\n", + "'f0': 0.0001032 [s^-1] (Coriolis parameter at the middle of the domain),\n", + "'n': 1.5 [nondim] (aspect ratio (n = 2 L_y / L_x)),\n", + "'rra': 6370000.0 [m] (earth radius),\n", + "'phi0_npi': 0.2777777777777778 [nondim] (latitude expressed in fraction of pi),\n", + "'deltap': 50000.0 [Pa] (pressure difference between the two atmospheric layers),\n", + "'Ha': 8500.0 [m] (Average height of the 500 hPa pressure level at midlatitude),\n", + "\n", + "Atmospheric Parameters:\n", + "'kd': 0.1 [nondim] (atmosphere bottom friction coefficient),\n", + "'kdp': 0.01 [nondim] (atmosphere internal friction coefficient),\n", + "'sigma': 0.2 [nondim] (static stability of the atmosphere),\n", + "\n", + "Atmospheric Temperature Parameters:\n", + "'gamma': 10000000.0 [J][m^-2][K^-1] (specific heat capacity of the atmosphere),\n", + "'C[1]': 99.0 [W][m^-2] (spectral component 1 of the short-wave radiation of the atmosphere),\n", + "'C[2]': 0.0 [W][m^-2] (spectral component 2 of the short-wave radiation of the atmosphere),\n", + "'C[3]': 0.0 [W][m^-2] (spectral component 3 of the short-wave radiation of the atmosphere),\n", + "'C[4]': 0.0 [W][m^-2] (spectral component 4 of the short-wave radiation of the atmosphere),\n", + "'C[5]': 0.0 [W][m^-2] (spectral component 5 of the short-wave radiation of the atmosphere),\n", + "'C[6]': 0.0 [W][m^-2] (spectral component 6 of the short-wave radiation of the atmosphere),\n", + "'C[7]': 0.0 [W][m^-2] (spectral component 7 of the short-wave radiation of the atmosphere),\n", + "'C[8]': 0.0 [W][m^-2] (spectral component 8 of the short-wave radiation of the atmosphere),\n", + "'C[9]': 0.0 [W][m^-2] (spectral component 9 of the short-wave radiation of the atmosphere),\n", + "'C[10]': 0.0 [W][m^-2] (spectral component 10 of the short-wave radiation of the atmosphere),\n", + "'eps': 0.76 (emissivity coefficient for the grey-body atmosphere),\n", + "'T0': 270.0 [K] (stationary solution for the 0-th order atmospheric temperature),\n", + "'sc': 1.0 (ratio of surface to atmosphere temperature),\n", + "'hlambda': 10.0 [W][m^-2][K^-1] (sensible+turbulent heat exchange between ocean/ground and atmosphere),\n", + "\n", + "Ground Parameters:\n", + "'hk[1]': 0.0 (spectral component 1 of the orography),\n", + "'hk[2]': 0.2 (spectral components 2 of the orography),\n", + "'hk[3]': 0.0 (spectral component 3 of the orography),\n", + "'hk[4]': 0.0 (spectral component 4 of the orography),\n", + "'hk[5]': 0.0 (spectral component 5 of the orography),\n", + "'hk[6]': 0.0 (spectral component 6 of the orography),\n", + "'hk[7]': 0.0 (spectral component 7 of the orography),\n", + "'hk[8]': 0.0 (spectral component 8 of the orography),\n", + "'hk[9]': 0.0 (spectral component 9 of the orography),\n", + "'hk[10]': 0.0 (spectral component 10 of the orography),\n", + "'orographic_basis': atmospheric,\n", + "\n", + "Ground Temperature Parameters:\n", + "'gamma': 16000000.0 [J][m^-2][K^-1] (specific heat capacity of the ground),\n", + "'C[1]': 300.0 [W][m^-2] (spectral component 1 of the short-wave radiation of the ground),\n", + "'C[2]': 0.0 [W][m^-2] (spectral component 2 of the short-wave radiation of the ground),\n", + "'C[3]': 0.0 [W][m^-2] (spectral component 3 of the short-wave radiation of the ground),\n", + "'C[4]': 0.0 [W][m^-2] (spectral component 4 of the short-wave radiation of the ground),\n", + "'C[5]': 0.0 [W][m^-2] (spectral component 5 of the short-wave radiation of the ground),\n", + "'C[6]': 0.0 [W][m^-2] (spectral component 6 of the short-wave radiation of the ground),\n", + "'C[7]': 0.0 [W][m^-2] (spectral component 7 of the short-wave radiation of the ground),\n", + "'C[8]': 0.0 [W][m^-2] (spectral component 8 of the short-wave radiation of the ground),\n", + "'C[9]': 0.0 [W][m^-2] (spectral component 9 of the short-wave radiation of the ground),\n", + "'C[10]': 0.0 [W][m^-2] (spectral component 10 of the short-wave radiation of the ground),\n", + "'T0': 285.0 [K] (stationary solution for the 0-th order oceanic temperature),\n", + "\n", + "\n" + ] + } + ], "source": [ "# Printing the model's parameters\n", "model_parameters.print_params()" @@ -191,9 +267,36 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 11, "metadata": {}, "outputs": [], + "source": [ + "from qgs.inner_products.symbolic import AtmosphericSymbolicInnerProducts" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [], + "source": [ + "aip = AtmosphericSymbolicInnerProducts(model_parameters)" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "CPU times: user 8.88 s, sys: 1.12 s, total: 10 s\n", + "Wall time: 8.55 s\n" + ] + } + ], "source": [ "%%time\n", "f, Df = create_tendencies(model_parameters)" diff --git a/notebooks/symbolic_outputs/symbolic_output_land_atmosphere-AUTO.ipynb b/notebooks/symbolic_outputs/symbolic_output_land_atmosphere-AUTO.ipynb new file mode 100644 index 0000000..06495d6 --- /dev/null +++ b/notebooks/symbolic_outputs/symbolic_output_land_atmosphere-AUTO.ipynb @@ -0,0 +1,889 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "037e82a6", + "metadata": {}, + "source": [ + "# Output of the symbolic tendencies: AUTO example " + ] + }, + { + "cell_type": "markdown", + "id": "9b0e76cb", + "metadata": {}, + "source": [ + "In this notebook, we show how to create the symbolic tendencies of the model and use them to perform continuation of solutions as parameters are varied, using the [AUTO-07p](https://github.com/auto-07p/auto-07p) continuation software. Symbolic tendencies here means that it is possible to make any parameter of the model appears in the tendencies equations.\n", + "\n", + "The present notebook will create the Fortran tendencies equations code and insert it into a given AUTO template.\n", + "This can then be used inside the notebook directly." + ] + }, + { + "cell_type": "markdown", + "id": "b8b7c180", + "metadata": {}, + "source": [ + "More details about the model used in this notebook can be found in the articles:\n", + "* Hamilton, O., Demaeyer, J., Vannitsem, S., & Crucifix, M. (2025). *Using Unstable Periodic Orbits to Understand Blocking Behaviour in a Low Order Land-Atmosphere Model*. Submitted to Chaos. [preprint](https://doi.org/10.48550/arXiv.2503.02808)\n", + "* Xavier, A. K., Demaeyer, J., & Vannitsem, S. (2024). *Variability and predictability of a reduced-order land–atmosphere coupled model.* Earth System Dynamics, **15**(4), 893-912. [doi:10.5194/esd-15-893-2024](https://doi.org/10.5194/esd-15-893-2024)\n", + "\n", + "or in the documentation. In particular, Hamilton et. al. (2025) used the symbolic tendencies here provided to compute periodic orbits, along with an automatic layer for AUTO called [auto-AUTO](https://github.com/Climdyn/auto-AUTO)." + ] + }, + { + "cell_type": "markdown", + "id": "98188e40", + "metadata": {}, + "source": [ + "> **To run this notebook, you need AUTO properly installed and configured !**\n", + ">\n", + "> **In general, it means that typing** `auto` **in a terminal starts the AUTO Python interface.**" + ] + }, + { + "cell_type": "markdown", + "id": "1c3c8250", + "metadata": {}, + "source": [ + "## Modules import" + ] + }, + { + "cell_type": "markdown", + "id": "cfb7a334", + "metadata": {}, + "source": [ + "First, setting the path and loading of some modules" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12f59a3f", + "metadata": {}, + "outputs": [], + "source": [ + "import sys, os\n", + "sys.path.extend([os.path.abspath('../')])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2d1499c4", + "metadata": {}, + "outputs": [], + "source": [ + "import glob" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4a29d8d6", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import matplotlib.pyplot as plt" + ] + }, + { + "cell_type": "markdown", + "id": "029a6f39", + "metadata": {}, + "source": [ + "Importing the model's modules" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7711e102", + "metadata": {}, + "outputs": [], + "source": [ + "from qgs.params.params import QgParams" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7e0d469c", + "metadata": {}, + "outputs": [], + "source": [ + "from qgs.functions.symbolic_tendencies import create_symbolic_equations" + ] + }, + { + "cell_type": "markdown", + "id": "c85c3ef7", + "metadata": {}, + "source": [ + "## Systems definition" + ] + }, + { + "cell_type": "markdown", + "id": "f5c74bfa", + "metadata": {}, + "source": [ + "Setting some model parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bf632fac", + "metadata": {}, + "outputs": [], + "source": [ + "model_parameters = QgParams({'phi0_npi': np.deg2rad(50.)/np.pi, 'n':1.3 }, dynamic_T=False)" + ] + }, + { + "cell_type": "markdown", + "id": "19326037", + "metadata": {}, + "source": [ + "and defining the spectral modes used by the model (they must be *symbolic*)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a2a8e7c1", + "metadata": {}, + "outputs": [], + "source": [ + "# Mode truncation at the wavenumber 2 in both x and y spatial coordinate for the atmosphere\n", + "model_parameters.set_atmospheric_channel_fourier_modes(2, 2, mode=\"symbolic\")\n", + "# Same modes for the ground temperature modes\n", + "model_parameters.set_ground_channel_fourier_modes(2, 2, mode=\"symbolic\")" + ] + }, + { + "cell_type": "markdown", + "id": "8bf273c3", + "metadata": {}, + "source": [ + "Completing the model parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3237dc76", + "metadata": {}, + "outputs": [], + "source": [ + "# Changing (increasing) the orography depth\n", + "model_parameters.ground_params.set_orography(0.2, 1)\n", + "# Setting the parameters of the heat transfer from the soil\n", + "model_parameters.gotemperature_params.set_params({'gamma': 1.6e7, 'T0': 300})\n", + "model_parameters.atemperature_params.set_params({ 'hlambda':10, 'T0': 290})\n", + "# Setting atmospheric parameters\n", + "model_parameters.atmospheric_params.set_params({'sigma': 0.2, 'kd': 0.085, 'kdp': 0.02})\n", + "\n", + "# Setting insolation \n", + "model_parameters.gotemperature_params.set_params({})" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "51907767", + "metadata": {}, + "outputs": [], + "source": [ + "C_g = 300\n", + "model_parameters.atemperature_params.set_insolation(0.4*C_g , 0)\n", + "\n", + "model_parameters.gotemperature_params.set_insolation(C_g , 0)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "979e2a7c", + "metadata": {}, + "outputs": [], + "source": [ + "# Printing the model's parameters\n", + "model_parameters.print_params()" + ] + }, + { + "cell_type": "markdown", + "id": "da46eccd", + "metadata": {}, + "source": [ + "## Creating AUTO files" + ] + }, + { + "cell_type": "markdown", + "id": "8ddc3286", + "metadata": {}, + "source": [ + "Calculating the tendencies in Fortran for AUTO as a function of the parameters $C_{{\\rm o},0}$, $C_{{\\rm a},0}$, $k_d$ and $k'_d$:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "af1b4747", + "metadata": {}, + "outputs": [], + "source": [ + "funcs, = create_symbolic_equations(model_parameters, continuation_variables=[model_parameters.gotemperature_params.C[0], model_parameters.atemperature_params.C[0], model_parameters.atmospheric_params.kd, model_parameters.atmospheric_params.kdp], language='auto')" + ] + }, + { + "cell_type": "markdown", + "id": "b21db3b5", + "metadata": {}, + "source": [ + "Let's inspect the output. First the AUTO `.f90` file:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2bdd3728", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "print(funcs[0])" + ] + }, + { + "cell_type": "markdown", + "id": "5cb3a9aa", + "metadata": {}, + "source": [ + "and then AUTO `c.` configuration file:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "73c0eb0a", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "print(funcs[1])" + ] + }, + { + "cell_type": "markdown", + "id": "bfe60c32", + "metadata": {}, + "source": [ + "We can now use both to write the AUTO files.\n", + "First we will modify the tendencies to force $C_{{\\rm a}, 1} = 0.4 C_{{\\rm g}, 1}$ , which is a standard assumption for these models, and reduces the number of parameters to define :" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "87f86011", + "metadata": {}, + "outputs": [], + "source": [ + "# splitting all the lines of the .f90 file\n", + "auto_eq_lines = funcs[0].split('\\n')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ec60bcec", + "metadata": {}, + "outputs": [], + "source": [ + "# forcing the change\n", + "for i, line in enumerate(auto_eq_lines):\n", + " if 'C_a1 = PAR(2)' in line:\n", + " auto_eq_lines[i] = '\\tC_a1 = 0.4*C_go1'\n", + " break" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "626c1862", + "metadata": {}, + "outputs": [], + "source": [ + "# gathering all the lines again in a single string\n", + "auto_eq = '\\n'.join(auto_eq_lines)" + ] + }, + { + "cell_type": "markdown", + "id": "b0d3fb95", + "metadata": {}, + "source": [ + "Taking care of the config file, changing some default settings:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "78c7ceba", + "metadata": {}, + "outputs": [], + "source": [ + "# splitting all the lines of the c. file\n", + "auto_config_lines = funcs[1].split('\\n')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "065e83b3", + "metadata": {}, + "outputs": [], + "source": [ + "# introducing some user defined points for AUTO\n", + "for i, line in enumerate(auto_config_lines):\n", + " if 'UZR' in line:\n", + " auto_config_lines[i] = \"UZR = {'C_go1': \" + str(list(np.arange(50.,375.,50.)))+\"}\"\n", + " break" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cc7b79fe", + "metadata": {}, + "outputs": [], + "source": [ + "# imposing that C_go1 is between 0. and 400. as stopping condition for AUTO\n", + "for i, line in enumerate(auto_config_lines):\n", + " if 'UZSTOP' in line:\n", + " auto_config_lines[i] = \"UZSTOP = {'C_go1': [0.,400.]}\"\n", + " break" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f7760092", + "metadata": {}, + "outputs": [], + "source": [ + "# gathering all the lines again in a single string\n", + "auto_config = '\\n'.join(auto_config_lines)" + ] + }, + { + "cell_type": "markdown", + "id": "5289260a", + "metadata": {}, + "source": [ + "and writing to files:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "80f05363", + "metadata": {}, + "outputs": [], + "source": [ + "model_name = 'qgs_land-atmosphere_auto'\n", + "with open(f'{model_name}.f90', 'w') as ff:\n", + " ff.write(auto_eq)\n", + " \n", + "with open(f'c.{model_name}', 'w') as ff:\n", + " ff.write(auto_config)" + ] + }, + { + "cell_type": "markdown", + "id": "3f219527", + "metadata": {}, + "source": [ + "## Defining some plotting functions" + ] + }, + { + "cell_type": "markdown", + "id": "5801aaa0", + "metadata": {}, + "source": [ + "to help us investigate the results later:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "29fe26ab", + "metadata": {}, + "outputs": [], + "source": [ + "def plot_branches(filename, variables=(0,1), ax=None, figsize=(10, 8), markersize=12., plot_kwargs=None, marker_kwargs=None, branch_indices='all', excluded_labels=('UZ', 'EP', 'No Label'), variables_name=None):\n", + " \n", + " if ax is None:\n", + " fig = plt.figure(figsize=figsize)\n", + " ax = fig.gca()\n", + " \n", + " if plot_kwargs is None:\n", + " plot_kwargs = dict()\n", + " \n", + " if marker_kwargs is None:\n", + " marker_kwargs = dict()\n", + " \n", + " pb_obj = parseB.parseB()\n", + " fb = open(filename, 'r')\n", + " pb_obj.read(fb)\n", + " \n", + " keys = list(pb_obj.branches[0].keys())\n", + " \n", + " if variables[0] in keys:\n", + " var1 = variables[0]\n", + " else:\n", + " try:\n", + " var1 = keys[variables[0]]\n", + " except:\n", + " var1 = keys[0]\n", + "\n", + " if variables[1] in keys:\n", + " var2 = variables[1]\n", + " else:\n", + " try:\n", + " var2 = keys[variables[1]]\n", + " except:\n", + " var2 = keys[1]\n", + "\n", + " if branch_indices == 'all':\n", + " branch_indices = range(len(pb_obj.branches))\n", + "\n", + " branch_num = list()\n", + " for i in branch_indices:\n", + " branch_dict = pb_obj.branches[i].todict()\n", + " branch_num.append(pb_obj.branches[i]['BR'])\n", + "\n", + " labels = list()\n", + " for j, coords in enumerate(zip(branch_dict[var1], branch_dict[var2])):\n", + " lab = pb_obj.branches[i].labels[j]\n", + " if not lab:\n", + " pass\n", + " else:\n", + " labels.append((coords, list(lab.keys())[0]))\n", + "\n", + " ax.plot(branch_dict[var1], branch_dict[var2], **plot_kwargs)\n", + " if excluded_labels != 'all':\n", + " for label in labels:\n", + " coords = label[0]\n", + " lab = label[1]\n", + " if lab not in excluded_labels:\n", + " ax.text(coords[0], coords[1], r'${\\bf '+ lab + r'}$', fontdict={'family':'sans-serif','size':markersize},va='center', ha='center', **marker_kwargs, clip_on=True)\n", + " \n", + " fb.close()\n", + " if variables_name is None:\n", + " ax.set_xlabel(var1)\n", + " ax.set_ylabel(var2)\n", + " else:\n", + " if isinstance(variables_name, dict):\n", + " ax.set_xlabel(variables_name[var1])\n", + " ax.set_ylabel(variables_name[var2])\n", + " else:\n", + " ax.set_xlabel(variables_name[0])\n", + " ax.set_ylabel(variables_name[1])\n", + " return ax, branch_num" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c1b0d5cb", + "metadata": {}, + "outputs": [], + "source": [ + "def plot_branches3d(filename, variables=(0,1,3), ax=None, figsize=(10, 8), markersize=12., plot_kwargs=None, marker_kwargs=None, branch_indices='all', excluded_labels=('UZ', 'EP', 'No Label'), variables_name=None):\n", + " \n", + " if ax is None:\n", + " fig = plt.figure(figsize=figsize)\n", + " ax = plt.subplot(projection='3d')\n", + " \n", + " if plot_kwargs is None:\n", + " plot_kwargs = dict()\n", + " \n", + " if marker_kwargs is None:\n", + " marker_kwargs = dict()\n", + " \n", + " pb_obj = parseB.parseB()\n", + " fb = open(filename, 'r')\n", + " pb_obj.read(fb)\n", + " \n", + " keys = list(pb_obj.branches[0].keys())\n", + " \n", + " if variables[0] in keys:\n", + " var1 = variables[0]\n", + " else:\n", + " try:\n", + " var1 = keys[variables[0]]\n", + " except:\n", + " var1 = keys[0]\n", + "\n", + " if variables[1] in keys:\n", + " var2 = variables[1]\n", + " else:\n", + " try:\n", + " var2 = keys[variables[1]]\n", + " except:\n", + " var2 = keys[1]\n", + " \n", + " if variables[2] in keys:\n", + " var3 = variables[2]\n", + " else:\n", + " try:\n", + " var3 = keys[variables[2]]\n", + " except:\n", + " var3 = keys[2]\n", + "\n", + "\n", + " if branch_indices == 'all':\n", + " branch_indices = range(len(pb_obj.branches))\n", + "\n", + " branch_num = list()\n", + " for i in branch_indices:\n", + " branch_dict = pb_obj.branches[i].todict()\n", + " branch_num.append(pb_obj.branches[i]['BR'])\n", + "\n", + " labels = list()\n", + " for j, coords in enumerate(zip(branch_dict[var1], branch_dict[var2], branch_dict[var3])):\n", + " lab = pb_obj.branches[i].labels[j]\n", + " if not lab:\n", + " pass\n", + " else:\n", + " labels.append((coords, list(lab.keys())[0]))\n", + "\n", + " ax.plot(branch_dict[var1], branch_dict[var2], branch_dict[var3], **plot_kwargs)\n", + " if excluded_labels != 'all':\n", + " for label in labels:\n", + " coords = label[0]\n", + " lab = label[1]\n", + " if lab not in excluded_labels:\n", + " ax.text(coords[0], coords[1], coords[2], r'${\\bf '+ lab + r'}$', fontdict={'family':'sans-serif','size':markersize},va='center', ha='center', **marker_kwargs, clip_on=True)\n", + " \n", + " fb.close()\n", + " if variables_name is None:\n", + " ax.set_xlabel(var1)\n", + " ax.set_ylabel(var2)\n", + " ax.set_zlabel(var3)\n", + " else:\n", + " if isinstance(variables_name, dict):\n", + " ax.set_xlabel(variables_name[var1])\n", + " ax.set_ylabel(variables_name[var2])\n", + " ax.set_zlabel(variables_name[var3])\n", + " else:\n", + " ax.set_xlabel(variables_name[0])\n", + " ax.set_ylabel(variables_name[1])\n", + " ax.set_zlabel(variables_name[2])\n", + " return ax, branch_num" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ff2a41a0", + "metadata": {}, + "outputs": [], + "source": [ + "def plot_branch_vs_others(branch_num, figsize=(10, 16), excluded_labels=('UZ', 'EP', 'No Label')):\n", + " \n", + " fig = plt.figure(figsize=figsize)\n", + " ax = plt.subplot(2,1,1)\n", + " ax3 = plt.subplot(2,1,2, projection='3d')\n", + " \n", + " \n", + " \n", + " fp = glob.glob('./b.fp*')\n", + " fp = [item for item in fp if '~' not in os.path.basename(item)]\n", + " fp = [item for item in fp if '_' not in os.path.basename(item)]\n", + " \n", + " for i in range(len(fp)-1,-1,-1):\n", + " \n", + " try:\n", + " num = int(fp[i][-2:])\n", + " except:\n", + " num = int(fp[i][-1])\n", + " \n", + " if num == branch_num:\n", + " plot_branches(fp[i], ax=ax, plot_kwargs={'color': 'tab:blue', 'zorder': 10.}, variables=(0, 1), variables_name=(r'$C_{\\rm o}$', r'$L_2$ norm'), excluded_labels=excluded_labels)\n", + " plot_branches3d(fp[i], ax=ax3, plot_kwargs={'color': 'tab:blue', 'zorder': 10.}, variables=(3, 0, 1), variables_name=(r'$\\psi_{{\\rm a}, 2}$', r'$C_{\\rm o}$', r'$L_2$ norm'), excluded_labels=excluded_labels)\n", + " else:\n", + " plot_branches(fp[i], ax=ax, plot_kwargs={'color': 'tab:orange'}, variables=(0, 1), variables_name=(r'$C_{\\rm o}$', r'$L_2$ norm'), excluded_labels=\"all\")\n", + " plot_branches3d(fp[i], ax=ax3, plot_kwargs={'color': 'tab:orange'}, variables=(3, 0, 1), variables_name=(r'$\\psi_{{\\rm a}, 2}$', r'$C_{\\rm o}$', r'$L_2$ norm'), excluded_labels=\"all\")\n", + "\n", + " " + ] + }, + { + "cell_type": "markdown", + "id": "2d4e97e3", + "metadata": {}, + "source": [ + "## AUTO analysis" + ] + }, + { + "cell_type": "markdown", + "id": "613647e4", + "metadata": {}, + "source": [ + "Initializing AUTO" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0b291604", + "metadata": {}, + "outputs": [], + "source": [ + "# Finding where AUTO is installed\n", + "auto_directory = os.environ['AUTO_DIR']" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a70a505f", + "metadata": {}, + "outputs": [], + "source": [ + "# Adding it to the path\n", + "sys.path.append(auto_directory + '/python/auto')\n", + "sys.path.append(auto_directory + '/python')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d1ddf964", + "metadata": {}, + "outputs": [], + "source": [ + "# Loading the needed AUTO Python interfaces\n", + "import AUTOCommands as ac\n", + "import AUTOclui as acl\n", + "import interactiveBindings as ib\n", + "import runAUTO as ra\n", + "import parseB, parseC, parseD, parseS, parseBandS" + ] + }, + { + "cell_type": "markdown", + "id": "a30354af", + "metadata": {}, + "source": [ + "Loading the model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d4f451e7", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"Loading model \"+model_name)" + ] + }, + { + "cell_type": "markdown", + "id": "539f2d00", + "metadata": {}, + "source": [ + "Starting a runner" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1240ace5", + "metadata": {}, + "outputs": [], + "source": [ + "runner = ra.runAUTO()\n", + "ac.load(model_name, runner=runner)" + ] + }, + { + "cell_type": "markdown", + "id": "672ef822", + "metadata": {}, + "source": [ + "Finding the first branch of fixed point" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fe0c0600", + "metadata": { + "scrolled": false + }, + "outputs": [], + "source": [ + "U_dic = {i+1: 0. for i in range(model_parameters.ndim)}\n", + "x = ac.run(model_name, U=U_dic, ICP=['C_go1'], PAR={3: model_parameters.atmospheric_params.kd, 4: model_parameters.atmospheric_params.kdp}, runner=runner)\n", + "ac.save(x,'fp1')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "375f5335", + "metadata": {}, + "outputs": [], + "source": [ + "plot_branch_vs_others(1)" + ] + }, + { + "cell_type": "markdown", + "id": "7f2075c3", + "metadata": {}, + "source": [ + "## Computing the periodic orbits (POs) out of the fixed point" + ] + }, + { + "cell_type": "markdown", + "id": "e00931b3", + "metadata": {}, + "source": [ + "Loading the branch and printing the summary" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b535d9be", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "r = ac.loadbd('fp1')\n", + "print(r.summary())" + ] + }, + { + "cell_type": "markdown", + "id": "12205b02", + "metadata": {}, + "source": [ + "Listing the Hopf bifurcation points" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12261493", + "metadata": {}, + "outputs": [], + "source": [ + "solutions_list = list()\n", + "ps_obj = parseS.parseS('./s.fp1')\n", + "pc_full_obj = parseC.parseC('c.'+model_name)\n", + "for i in range(len(ps_obj)):\n", + " s = ps_obj[i].load(constants=pc_full_obj)\n", + " if s['TY'] == 'HB':\n", + " solutions_list.append(s)\n", + " \n", + "\n", + "# reversing to get it in Co increasing order\n", + "solutions_list = solutions_list[::-1]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f3932386", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "solutions_list" + ] + }, + { + "cell_type": "markdown", + "id": "91199f37", + "metadata": {}, + "source": [ + "### Computing and plotting the second Hopf bifurcation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a2604fb2", + "metadata": { + "scrolled": false + }, + "outputs": [], + "source": [ + "s = solutions_list[1]\n", + "rh=ac.run(s,ICP=['C_go1', 'T'], IPS=2, NTST=400, runner=runner)\n", + "ac.save(rh, 'fp1_hp1')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0593b6a8", + "metadata": { + "scrolled": false + }, + "outputs": [], + "source": [ + "ax, _ = plot_branches('./b.fp1_hp1', variables=(0, 3))\n", + "plot_branches('./b.fp1', ax=ax, variables=(0, 3))" + ] + }, + { + "cell_type": "markdown", + "id": "9b8f7205", + "metadata": {}, + "source": [ + "Other fixed point and periodic orbit branches can of course be computed in the same way." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.8.16" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/symbolic_outputs/symbolic_output_land_atmosphere.ipynb b/notebooks/symbolic_outputs/symbolic_output_land_atmosphere.ipynb new file mode 100644 index 0000000..822fe14 --- /dev/null +++ b/notebooks/symbolic_outputs/symbolic_output_land_atmosphere.ipynb @@ -0,0 +1,507 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "037e82a6", + "metadata": {}, + "source": [ + "# Output of the symbolic tendencies: Land-Atmosphere model example " + ] + }, + { + "cell_type": "markdown", + "id": "9b0e76cb", + "metadata": {}, + "source": [ + "In this notebook, we show how to create the symbolic tendencies of the model. Symbolic tendencies here means that it is possible to make any parameter of the model appears in the tendencies equations.\n", + "\n", + "This can be done in several languages (Python, Julia, Fortran), but here, we are going to use Python." + ] + }, + { + "cell_type": "markdown", + "id": "b8b7c180", + "metadata": {}, + "source": [ + "More details about the model used in this notebook can be found in the articles:\n", + "* Hamilton, O., Demaeyer, J., Vannitsem, S., & Crucifix, M. (2025). *Using Unstable Periodic Orbits to Understand Blocking Behaviour in a Low Order Land-Atmosphere Model*. Submitted to Chaos. [preprint](https://doi.org/10.48550/arXiv.2503.02808)\n", + "* Xavier, A. K., Demaeyer, J., & Vannitsem, S. (2024). *Variability and predictability of a reduced-order land–atmosphere coupled model.* Earth System Dynamics, **15**(4), 893-912. [doi:10.5194/esd-15-893-2024](https://doi.org/10.5194/esd-15-893-2024)\n", + "\n", + "or in the documentation." + ] + }, + { + "cell_type": "markdown", + "id": "1c3c8250", + "metadata": {}, + "source": [ + "## Modules import" + ] + }, + { + "cell_type": "markdown", + "id": "cfb7a334", + "metadata": {}, + "source": [ + "First, setting the path and loading of some modules" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12f59a3f", + "metadata": {}, + "outputs": [], + "source": [ + "import sys, os\n", + "sys.path.extend([os.path.abspath('../')])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4a29d8d6", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import matplotlib.pyplot as plt" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1a13ae55", + "metadata": {}, + "outputs": [], + "source": [ + "from numba import njit" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2468708a", + "metadata": {}, + "outputs": [], + "source": [ + "from scipy.integrate import solve_ivp" + ] + }, + { + "cell_type": "markdown", + "id": "029a6f39", + "metadata": {}, + "source": [ + "Importing the model's modules" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7711e102", + "metadata": {}, + "outputs": [], + "source": [ + "from qgs.params.params import QgParams" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7e0d469c", + "metadata": {}, + "outputs": [], + "source": [ + "from qgs.functions.symbolic_tendencies import create_symbolic_equations\n", + "from qgs.functions.tendencies import create_tendencies" + ] + }, + { + "cell_type": "markdown", + "id": "c85c3ef7", + "metadata": {}, + "source": [ + "## Systems definition" + ] + }, + { + "cell_type": "markdown", + "id": "0a2b894b", + "metadata": {}, + "source": [ + "General parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d73ac5f8", + "metadata": {}, + "outputs": [], + "source": [ + "# Time parameters\n", + "dt = 0.1\n" + ] + }, + { + "cell_type": "markdown", + "id": "f5c74bfa", + "metadata": {}, + "source": [ + "Setting some model parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bf632fac", + "metadata": {}, + "outputs": [], + "source": [ + "model_parameters = QgParams({'phi0_npi': np.deg2rad(50.)/np.pi, 'n':1.3 }, dynamic_T=False)" + ] + }, + { + "cell_type": "markdown", + "id": "19326037", + "metadata": {}, + "source": [ + "and defining the spectral modes used by the model (they must be *symbolic*)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a2a8e7c1", + "metadata": {}, + "outputs": [], + "source": [ + "# Mode truncation at the wavenumber 2 in both x and y spatial coordinate for the atmosphere\n", + "model_parameters.set_atmospheric_channel_fourier_modes(2, 2, mode=\"symbolic\")\n", + "# Same modes for the ground temperature modes\n", + "model_parameters.set_ground_channel_fourier_modes(2, 2, mode=\"symbolic\")" + ] + }, + { + "cell_type": "markdown", + "id": "8bf273c3", + "metadata": {}, + "source": [ + "Completing the model parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3237dc76", + "metadata": {}, + "outputs": [], + "source": [ + "# Changing (increasing) the orography depth\n", + "model_parameters.ground_params.set_orography(0.2, 1)\n", + "# Setting the parameters of the heat transfer from the soil\n", + "model_parameters.gotemperature_params.set_params({'gamma': 1.6e7, 'T0': 300})\n", + "model_parameters.atemperature_params.set_params({ 'hlambda':10, 'T0': 290})\n", + "# Setting atmospheric parameters\n", + "model_parameters.atmospheric_params.set_params({'sigma': 0.2, 'kd': 0.085, 'kdp': 0.02})\n", + "\n", + "# Setting insolation \n", + "model_parameters.gotemperature_params.set_params({})" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "51907767", + "metadata": {}, + "outputs": [], + "source": [ + "C_g = 300\n", + "model_parameters.atemperature_params.set_insolation(0.4*C_g , 0)\n", + "\n", + "model_parameters.gotemperature_params.set_insolation(C_g , 0)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "979e2a7c", + "metadata": {}, + "outputs": [], + "source": [ + "# Printing the model's parameters\n", + "model_parameters.print_params()" + ] + }, + { + "cell_type": "markdown", + "id": "d43d341b", + "metadata": {}, + "source": [ + "## Outputting the model equations" + ] + }, + { + "cell_type": "markdown", + "id": "8ddc3286", + "metadata": {}, + "source": [ + "Calculating the tendencies in Python as a function of the parameters $C_{{\\rm g},0}$, $C_{{\\rm a},0}$, $k_d$ and $k'_d$:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9479fcb0", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "%%time\n", + "funcs, = create_symbolic_equations(model_parameters, continuation_variables=[model_parameters.gotemperature_params.C[0], model_parameters.atemperature_params.C[0], model_parameters.atmospheric_params.kd, model_parameters.atmospheric_params.kdp], language='python')" + ] + }, + { + "cell_type": "markdown", + "id": "c61d613a", + "metadata": {}, + "source": [ + "Let's inspect the output:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bdb4d33f", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "print(funcs)" + ] + }, + { + "cell_type": "markdown", + "id": "23ccb10c", + "metadata": {}, + "source": [ + "Note that the tendencies have been already formatted as a [Numba](https://numba.pydata.org/) function, but it is easy to extract the equations for any other kind of accelerator or simply to produce pure Python code." + ] + }, + { + "cell_type": "markdown", + "id": "da0c63c1", + "metadata": {}, + "source": [ + "It is now easy to get the function into operation:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a8fea890", + "metadata": {}, + "outputs": [], + "source": [ + "exec(funcs)" + ] + }, + { + "cell_type": "markdown", + "id": "992a6243", + "metadata": {}, + "source": [ + "and" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2ae15534", + "metadata": {}, + "outputs": [], + "source": [ + "f(0.,np.zeros(model_parameters.ndim), 300., 120., 0.085, 0.02)" + ] + }, + { + "cell_type": "markdown", + "id": "0579478b", + "metadata": {}, + "source": [ + "## Comparing with numerical results" + ] + }, + { + "cell_type": "markdown", + "id": "b9be8b1d", + "metadata": {}, + "source": [ + "We can check that the symbolic (parameters dependent) equations are the same as the `qgs` numerical ones (with the same values of the parameters):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f38b345f", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "f_num, Df = create_tendencies(model_parameters)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "46d33700", + "metadata": {}, + "outputs": [], + "source": [ + "f_num(0., np.zeros(model_parameters.ndim))" + ] + }, + { + "cell_type": "markdown", + "id": "ea417bb2", + "metadata": {}, + "source": [ + "In addition, one can easily compare the obtained attractors:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9a0f5941", + "metadata": {}, + "outputs": [], + "source": [ + "# IC calculated from a long transient\n", + "ic = np.array([0.05055959, -0.01639403, -0.01440781, -0.01846523, -0.01352099,\n", + " 0.011685 , -0.00201673, -0.02030682, 0.03923588, -0.02229535,\n", + " 0.0586372 , -0.01805569, -0.01264252, -0.0103574 , -0.00618456,\n", + " 0.01159318, -0.00478694, -0.00782509, 0.01066059, -0.01552667,\n", + " 0.30718325, -0.03247899, -0.04512935, -0.00078786, -0.00067468,\n", + " 0.00183836, 0.00068025, 0.00215424, -0.00322845, -0.00186392])\n", + "\n", + "# Actual integration\n", + "traj = solve_ivp(f, (0., 100000.), ic, t_eval=np.arange(0, 100000., dt), args=(300., 120., 0.085, 0.02))\n", + "traj_num = solve_ivp(f_num, (0., 100000.), ic, t_eval=np.arange(0, 100000., dt))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2f892676", + "metadata": {}, + "outputs": [], + "source": [ + "varx = 2\n", + "vary = 1\n", + "plt.figure(figsize=(12, 10))\n", + "\n", + "plt.plot(traj.y[varx], traj.y[vary], marker='o', ms=0.03, ls='', label='Symbolic tendencies')\n", + "plt.plot(traj_num.y[varx], traj_num.y[vary], marker='o', ms=0.03, ls='', label='Fully numerical tendencies')\n", + "\n", + "\n", + "plt.xlabel('$'+model_parameters.latex_var_string[varx]+'$')\n", + "plt.ylabel('$'+model_parameters.latex_var_string[vary]+'$');\n", + "# plt.legend()" + ] + }, + { + "cell_type": "markdown", + "id": "91781f9f", + "metadata": {}, + "source": [ + "Fully numerical tendencies attractor is in orange while the symbolic tendencies on is in blue" + ] + }, + { + "cell_type": "markdown", + "id": "f84b7484", + "metadata": {}, + "source": [ + "## Varying the parameters" + ] + }, + { + "cell_type": "markdown", + "id": "6367bee4", + "metadata": {}, + "source": [ + "The obvious possibilities given by the symbolic tendencies are to allow users to easily perform sensitivity analysis:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "291b2342", + "metadata": {}, + "outputs": [], + "source": [ + "# let's start with 4 different values of the friction k_d\n", + "kd_list = [0.06, 0.085, 0.095, 0.1, 0.105, 0.12]\n", + "\n", + "# let's compute the attractor for each\n", + "\n", + "attractor_list = list()\n", + "\n", + "for kd in kd_list:\n", + " attractor_list.append(solve_ivp(f, (0., 100000.), ic, t_eval=np.arange(0, 100000., dt), args=(300., 120., kd, 0.02)))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e068b639", + "metadata": { + "scrolled": false + }, + "outputs": [], + "source": [ + "varx = 2\n", + "vary = 1\n", + "plt.figure(figsize=(18, 6))\n", + "\n", + "k=1\n", + "for kd, attractor in zip(kd_list, attractor_list):\n", + " plt.subplot(2, 3, k)\n", + " plt.plot(attractor.y[varx], attractor.y[vary], marker='o', ms=0.03, ls='', label=f'$k_d$ = {kd}')\n", + " plt.xlabel('$'+model_parameters.latex_var_string[varx]+'$')\n", + " plt.ylabel('$'+model_parameters.latex_var_string[vary]+'$');\n", + "\n", + " plt.legend()\n", + " k+=1\n", + "\n", + "\n", + "\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.8.16" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/qgs/basis/base.py b/qgs/basis/base.py index 73f2082..7999ef7 100644 --- a/qgs/basis/base.py +++ b/qgs/basis/base.py @@ -26,7 +26,7 @@ import sys from abc import ABC -from sympy import symbols, lambdify, diff +from sympy import Symbol, symbols, lambdify, diff class Basis(ABC): @@ -219,7 +219,8 @@ def num_functions(self): basis = SymbolicBasis() x, y = symbols('x y') # x and y coordinates on the model's spatial domain - n, al = symbols('n al') # aspect ratio and alpha coefficients + al = symbols('al') # aspect ratio and alpha coefficients + n = Symbol('n', positive=True) for i in range(1, 3): for j in range(1, 3): basis.append(2 * exp(- al * x) * sin(j * n * x / 2) * sin(i * y)) diff --git a/qgs/basis/fourier.py b/qgs/basis/fourier.py index f645d91..1a14794 100644 --- a/qgs/basis/fourier.py +++ b/qgs/basis/fourier.py @@ -21,7 +21,7 @@ from sympy import symbols, sin, cos, sqrt _x, _y = symbols('x y') -_n = symbols('n', real=True, nonnegative=True) +_n = symbols('n', positive=True) class ChannelFourierBasis(SymbolicBasis): diff --git a/qgs/functions/sparse_mul.py b/qgs/functions/sparse_mul.py index cb2707e..22ada30 100644 --- a/qgs/functions/sparse_mul.py +++ b/qgs/functions/sparse_mul.py @@ -9,10 +9,11 @@ import numpy as np from numba import njit + @njit def sparse_mul2(coo, value, vec): """Sparse multiplication of a tensor with one vector: - :math:`A_{i,j} = {\displaystyle \sum_{k=0}^{\mathrm{ndim}}} \, \mathcal{T}_{i,j,k} \, a_k` + :math:`A_{i,j} = {\\displaystyle \\sum_{k=0}^{\\mathrm{ndim}}} \\, \\mathcal{T}_{i,j,k} \\, a_k` Warnings -------- @@ -47,7 +48,7 @@ def sparse_mul2(coo, value, vec): @njit def sparse_mul3(coo, value, vec_a, vec_b): """Sparse multiplication of a tensor with two vectors: - :math:`v_i = {\displaystyle \sum_{j,k=0}^{\mathrm{ndim}}} \, \mathcal{T}_{i,j,k} \, a_j \, b_k` + :math:`v_i = {\\displaystyle \\sum_{j,k=0}^{\\mathrm{ndim}}} \\, \\mathcal{T}_{i,j,k} \\, a_j \\, b_k` Warnings -------- @@ -79,10 +80,11 @@ def sparse_mul3(coo, value, vec_a, vec_b): res[0] = 1. return res + @njit def sparse_mul4(coo, value, vec_a, vec_b, vec_c): """Sparse multiplication of a rank-5 tensor with three vectors: - :math:`A_{i, j} = {\displaystyle \sum_{k,l,m=0}^{\mathrm{ndim}}} \, \mathcal{T}_{i,j,k,l, m} \, a_k \, b_l \, c_m` + :math:`A_{i, j} = {\\displaystyle \\sum_{k,l,m=0}^{\\mathrm{ndim}}} \\, \\mathcal{T}_{i,j,k,l, m} \\, a_k \\, b_l \\, c_m` Warnings -------- @@ -119,7 +121,7 @@ def sparse_mul4(coo, value, vec_a, vec_b, vec_c): @njit def sparse_mul5(coo, value, vec_a, vec_b, vec_c, vec_d): """Sparse multiplication of a rank-5 tensor with four vectors: - :math:`v_i = {\displaystyle \sum_{j,k,l,m=0}^{\mathrm{ndim}}} \, \mathcal{T}_{i,j,k,l,m} \, a_j \, b_k \, c_l \, d_m` + :math:`v_i = {\\displaystyle \\sum_{j,k,l,m=0}^{\\mathrm{ndim}}} \\, \\mathcal{T}_{i,j,k,l,m} \\, a_j \\, b_k \\, c_l \\, d_m` Warnings -------- diff --git a/qgs/functions/symbolic_mul.py b/qgs/functions/symbolic_mul.py new file mode 100644 index 0000000..3707ea4 --- /dev/null +++ b/qgs/functions/symbolic_mul.py @@ -0,0 +1,199 @@ +""" + Symbolic matrix operation module + ================================ + + This module supports operations and functions on the symbolic sparse tensors defined in + the :class:`~.tensors.symbolic_qgtensor.SymbolicQgsTensor` objects. + +""" + +from sympy import tensorproduct, tensorcontraction + + +def add_to_dict(dic, loc, value): + """Adds `value` to dictionary `dic`, with the dictionary key of `loc`. + If the dictionary did not have a key of `loc` before, a new key is made. + + # Jonathan: Add parameters descriptions + """ + + try: + dic[loc] += value + except: + dic[loc] = value + return dic + + +def symbolic_tensordot(a, b, axes=2): + """Compute tensor dot product along specified axes of two sympy symbolic arrays + + This is based on `Numpy`_ :meth:`~numpy.tensordot` . + + Parameters + ---------- + a, b: ~sympy.tensor.array.DenseNDimArray or ~sympy.tensor.array.SparseNDimArray + Arrays to take the dot product of. + + axes: int + Sum over the last `axes` axes of `a` and the first `axes` axes + of `b` in order. The sizes of the corresponding axes must match. + + Returns + ------- + output: sympy tensor + The tensor dot product of the input. + + .. _Numpy: https://numpy.org/ + + """ + as_ = a.shape + nda = len(as_) + + a_com = [nda+i for i in range(-axes, 0)] + b_com = [nda+i for i in range(axes)] + sum_cols = tuple(a_com + b_com) + + prod = tensorproduct(a, b) + + return tensorcontraction(prod, sum_cols) + + +def symbolic_sparse_mult2(dic, vec_a): + """Symbolic multiplication of a tensor with one vector: + :math:`A_{i,j} = {\\displaystyle \\sum_{k=0}^{\\mathrm{ndim}}} \\, \\mathcal{T}_{i,j,k} \\, a_k` + + Parameters + ---------- + dic: dict(~sympy.core.symbol.Symbol) + A dictionary whose keys are the coordinates of the tensor, and the dictionary values are the values of the + tensor. + vec_a: list(~sympy.core.symbol.Symbol) + The list :math:`a_k` to contract the tensor with entries of Sympy Symbols. + Must be of shape (:attr:`~.params.QgParams.ndim` + 1,). + + Returns + ------- + res: dict(~sympy.core.symbol.Symbol) + The matrix :math:`A_{i,j}`, of shape (:attr:`~.params.QgParams.ndim` + 1, :attr:`~.params.QgParams.ndim` + 1), + contained in a dictionary, where the keys are the tensor coordinates, and the values are the tensor values. + """ + res = dict() + + for key in dic.keys(): + coo1, coo2, coo3 = key + if coo1 > 0 and coo2 > 0: + val = vec_a[coo3] * dic[key] + res = add_to_dict(res, (coo1, coo2), val) + + return res + + +def symbolic_sparse_mult3(dic, vec_a, vec_b): + """Symbolic multiplication of a tensor with two vectors: + :math:`v_i = {\\displaystyle \\sum_{j,k=0}^{\\mathrm{ndim}}} \\, \\mathcal{T}_{i,j,k} \\, a_j \\, b_k` + + Parameters + ---------- + dic: dict(~sympy.core.symbol.Symbol) + A dictionary whose keys are the coordinates of the tensor, and the dictionary values are the values of the + tensor. + vec_a: list(~sympy.core.symbol.Symbol) + The list :math:`a_j` to contract the tensor with entries of Sympy Symbols. + Must be of shape (:attr:`~.params.QgParams.ndim` + 1,). + vec_b: list(~sympy.core.symbol.Symbol) + The list :math:`b_k` to contract the tensor with entries of Sympy Symbols. + Must be of shape (:attr:`~.params.QgParams.ndim` + 1,). + + + Returns + ------- + res: dict(~sympy.core.symbol.Symbol) + The vector :math:`v_i`, of shape (:attr:`~.params.QgParams.ndim` + 1,), contained in a dictionary, where + the keys are the tensor coordinates, and the values are the tensor values. + + """ + res = dict() + + for key in dic.keys(): + coo1, coo2, coo3 = key + if coo1 > 0: + val = vec_a[coo2] * vec_b[coo3] * dic[key] + res = add_to_dict(res, coo1, val) + + return res + + +def symbolic_sparse_mult4(dic, vec_a, vec_b, vec_c): + """Symbolic multiplication of a rank-5 tensor with three vectors: + :math:`A_{i, j} = {\\displaystyle \\sum_{k,l,m=0}^{\\mathrm{ndim}}} \\, \\mathcal{T}_{i,j,k,l, m} \\, a_k \\, b_l \\, c_m` + + + Parameters + ---------- + dic: dict(~sympy.core.symbol.Symbol) + A dictionary where they keys are a tuple of 5 elements which are the coordinates of the tensor values, + which are contained in the dictionary values. + vec_a: list(~sympy.core.symbol.Symbol) + The list :math:`a_j` to contract the tensor with entries of Sympy Symbols. + Must be of shape (:attr:`~.params.QgParams.ndim` + 1,). + vec_b: list(~sympy.core.symbol.Symbol) + The list :math:`b_k` to contract the tensor with entries of Sympy Symbols. + Must be of shape (:attr:`~.params.QgParams.ndim` + 1,). + vec_c: list(~sympy.core.symbol.Symbol) + The list :math:`c_l` to contract the tensor with entries of Sympy Symbols. + Must be of shape (:attr:`~.params.QgParams.ndim` + 1,). + + Returns + ------- + res: dict(~sympy.core.symbol.Symbol) + The matrix :math:`A_{i, j}`, of shape (:attr:`~.params.QgParams.ndim` + 1, :attr:`~.params.QgParams.ndim` + 1), + contained in a dictionary, where the keys are the tensor coordinates, and the values are the tensor values. + """ + res = dict() + + for key in dic.keys(): + coo1, coo2, coo3, coo4, coo5 = key + if coo1 > 0 and coo2 > 0: + val = vec_a[coo3] * vec_b[coo4] * vec_c[coo5] * dic[key] + res = add_to_dict(res, (coo1, coo2), val) + + return res + + +def symbolic_sparse_mult5(dic, vec_a, vec_b, vec_c, vec_d): + """Symbolic multiplication of a rank-5 tensor with four vectors: + :math:`v_i = {\\displaystyle \\sum_{j,k,l,m=0}^{\\mathrm{ndim}}} \\, \\mathcal{T}_{i,j,k,l,m} \\, a_j \\, b_k \\, c_l \\, d_m` + + Parameters + ---------- + dic: dict(~sympy.core.symbol.Symbol) + A dictionary whose keys are the coordinates of the tensor, and the dictionary values are the values of the + tensor. + vec_a: list(~sympy.core.symbol.Symbol) + The list :math:`a_j` to contract the tensor with entries of Sympy Symbols. + Must be of shape (:attr:`~.params.QgParams.ndim` + 1,). + vec_b: list(~sympy.core.symbol.Symbol) + The list :math:`b_k` to contract the tensor with entries of Sympy Symbols. + Must be of shape (:attr:`~.params.QgParams.ndim` + 1,). + vec_c: list(~sympy.core.symbol.Symbol) + The list :math:`c_l` to contract the tensor with entries of Sympy Symbols. + Must be of shape (:attr:`~.params.QgParams.ndim` + 1,). + vec_d: list(~sympy.core.symbol.Symbol) + The list :math:`d_m` to contract the tensor with entries of Sympy Symbols. + Must be of shape (:attr:`~.params.QgParams.ndim` + 1,). + + Returns + ------- + res: dict(~sympy.core.symbol.Symbol) + The vector :math:`v_i`, of shape (:attr:`~.params.QgParams.ndim` + 1,), contained in a dictionary, + where the keys are the tensor coordinates, and the values are the tensor values. + """ + res = dict() + + for key in dic.keys(): + coo1, coo2, coo3, coo4, coo5 = key + if coo1 > 0: + val = vec_a[coo2] * vec_b[coo3] * vec_c[coo4] * vec_d[coo5] * dic[key] + res = add_to_dict(res, coo1, val) + + return res diff --git a/qgs/functions/symbolic_tendencies.py b/qgs/functions/symbolic_tendencies.py new file mode 100644 index 0000000..4fd9160 --- /dev/null +++ b/qgs/functions/symbolic_tendencies.py @@ -0,0 +1,1034 @@ +""" + Symbolic tendencies module + ========================== + + This module provides functions to create a symbolic representation of the tendencies functions of the model + in various languages and for various external software. + +""" + +import numpy as np +from sympy import Symbol, lambdify + +import warnings +from qgs.functions.symbolic_mul import symbolic_sparse_mult2, symbolic_sparse_mult3, symbolic_sparse_mult4, \ + symbolic_sparse_mult5 +from qgs.inner_products.symbolic import AtmosphericSymbolicInnerProducts, OceanicSymbolicInnerProducts, \ + GroundSymbolicInnerProducts +from qgs.tensors.symbolic_qgtensor import SymbolicQgsTensor, SymbolicQgsTensorDynamicT, SymbolicQgsTensorT4 + +python_lang_translation = { + 'sqrt': 'math.sqrt', + 'pi': 'math.pi', + 'lambda': 'lmda' # Remove conflict for lambda function in python +} + +fortran_lang_translation = { + 'conjugate': 'CONJG', + 'epsilon': 'eps' # Remove conflict for EPSILON function in fortran + # TODO: may need to add variable for pi +} + +julia_lang_translation = { + '**': '^', + 'pi': 'pi()', + 'conjugate': 'conj' +} + +mathematica_lang_translation = { + '**': '^' +} + + +def create_symbolic_equations(params, atm_ip=None, ocn_ip=None, gnd_ip=None, continuation_variables=None, + language='python', return_inner_products=False, return_jacobian=False, + return_symbolic_eqs=False, return_symbolic_qgtensor=False): + """Function to output the raw symbolic functions of the qgs model. + + Parameters + ---------- + params: QgParams + The parameters fully specifying the model configuration. + atm_ip: SymbolicAtmosphericInnerProducts, optional + Allows for stored inner products to be input. + ocn_ip: SymbolicOceanInnerProducts, optional + Allows for stored inner products to be input. + gnd_ip: SymbolicGroundInnerProducts, optional + Allows for stored inner products to be input. + continuation_variables: Iterable(Parameter, ScalingParameter, ParametersArray) + The variables to not substitute and to leave in the equations, if `None` all variables are substituted. + language: str + Options for the output language syntax: 'python', 'julia', 'fortran', 'auto', 'mathematica'. + Default to 'python'. + return_inner_products: bool + If `True`, return the inner products of the model. Default to `False`. + return_jacobian: bool + If `True`, return the Jacobian of the model. Default to `False`. + return_symbolic_eqs: bool + If `True`, return the substituted symbolic equations. + return_symbolic_qgtensor: bool + If `True`, return the symbolic tendencies tensor of the model. Default to `False`. + + Returns + ------- + funcs: str + The substituted functions in the language syntax specified, as a string. + dict_eq_simplified: dict(~sympy.core.expr.Expr) + Dictionary of the substituted Jacobian matrix. + inner_products: tuple(SymbolicAtmosphericInnerProducts, SymbolicOceanicInnerProducts, SymbolicGroundInnerProducts) + If `return_inner_products` is `True`, the inner products of the system. + eq_simplified: dict(~sympy.core.expr.Expr) + If `return_symbolic_eqs` is `True`, dictionary of the model tendencies symbolic functions. + agotensor: SymbolicQgsTensor + If `return_symbolic_qgtensor` is `True`, the symbolic tendencies tensor of the system. + + """ + make_ip_subs = True + + if continuation_variables is None: + make_ip_subs = False # TODO: check this !!! For me it should be True because if no cv then all subs are done. + else: + for cv in continuation_variables: + try: + if params.scale_params.n == cv: + make_ip_subs = False + except: + pass + + if not make_ip_subs: + warnings.warn("Calculating inner products symbolically, as the variable 'n' has been specified as a variable, " + "this may take a while.") + + if params.atmospheric_basis is not None: + if atm_ip is None: + aip = AtmosphericSymbolicInnerProducts(params, return_symbolic=True, make_substitution=make_ip_subs) + else: + aip = atm_ip + else: + aip = None + + if params.oceanic_basis is not None: + if ocn_ip is None: + oip = OceanicSymbolicInnerProducts(params, return_symbolic=True, make_substitution=make_ip_subs) + else: + oip = ocn_ip + else: + oip = None + + if params.ground_basis is not None: + if gnd_ip is None: + gip = GroundSymbolicInnerProducts(params, return_symbolic=True, make_substitution=make_ip_subs) + else: + gip = gnd_ip + else: + gip = None + + if aip is not None and oip is not None: + if not aip.connected_to_ocean: + aip.connect_to_ocean(oip) + elif aip is not None and gip is not None: + if not aip.connected_to_ground: + aip.connect_to_ground(gip) + + if params.T4: + agotensor = SymbolicQgsTensorT4(params, aip, oip, gip) + elif params.dynamic_T: + agotensor = SymbolicQgsTensorDynamicT(params, aip, oip, gip) + else: + agotensor = SymbolicQgsTensor(params, aip, oip, gip) + + xx = list() + xx.append(1) + + for i in range(1, params.ndim+1): + xx.append(Symbol('U_'+str(i))) + + if params.dynamic_T: + eq = symbolic_sparse_mult5(agotensor.sub_tensor(continuation_variables=continuation_variables), xx, xx, xx, xx) + if return_jacobian: + dict_eq = symbolic_sparse_mult4(agotensor.sub_tensor( + agotensor.jac_dic, continuation_variables=continuation_variables), xx, xx, xx) + + else: + eq = symbolic_sparse_mult3(agotensor.sub_tensor(continuation_variables=continuation_variables), xx, xx) + if return_jacobian: + dict_eq = symbolic_sparse_mult2(agotensor.sub_tensor( + agotensor.jac_dic, continuation_variables=continuation_variables), xx) + + eq_simplified = dict() + dict_eq_simplified = dict() + + if continuation_variables is None: + # Simplifying at this step is slow + # This only needs to be used if no substitutions are being made + for i in range(1, params.ndim+1): + eq_simplified[i] = eq[i].simplify() + if return_jacobian: + for j in range(1, params.ndim+1): + if (i, j) in dict_eq: + dict_eq_simplified[(i, j)] = dict_eq[(i, j)].simplify() + + else: + eq_simplified = eq + if return_jacobian: + dict_eq_simplified = dict_eq + + func = equation_as_function(equations=eq_simplified, params=params, language=language, + continuation_variables=continuation_variables) + + if return_jacobian: + func_jac = jacobian_as_function(equations=dict_eq_simplified, params=params, language=language, + continuation_variables=continuation_variables) + + ret = list() + ret.append(func) + if return_jacobian: + ret.append(func_jac) + if return_inner_products: + ret.append((aip, oip, gip)) + if return_symbolic_eqs: + ret.append(eq_simplified) + if return_symbolic_qgtensor: + ret.append(agotensor) + return ret + + +def translate_equations(equations, language='python'): + """Function to output the model equations as a string in the specified language syntax. + + Parameters + ---------- + equations: dict(string), list, string + Dictionary, list, or string of the symbolic model equations. + language: string + Language syntax that the equations are returned in. Options are: + - `python` + - `fortran` + - `julia` + - `auto` + - `mathematica` + Default to `python`. + + Returns + ------- + str_eq: dict(str) + Dictionary of strings of the model equations. + """ + + if language == 'python': + translator = python_lang_translation + + if language == 'julia': + translator = julia_lang_translation + + if language == 'fortran': + translator = fortran_lang_translation + + # translate mathematical operations + if isinstance(equations, dict): + str_eq = dict() + for key in equations.keys(): + temp_str = equations[key] + for k in translator.keys(): + temp_str = temp_str.replace(k, translator[k]) + str_eq[key] = temp_str + elif isinstance(equations, list): + str_eq = list() + for eq in equations: + for k in translator.keys(): + eq = eq.replace(k, translator[k]) + str_eq.append(eq) + elif isinstance(equations, str): + str_eq = equations + for k in translator.keys(): + str_eq = str_eq.replace(k, translator[k]) + else: + raise warnings.warn("Expected a dict, list, or string input") + + return str_eq + + +def format_equations(equations, params, save_loc=None, language='python', print_equations=False): + """Function formats the equations, in the programming language specified, and saves the equations to the specified + location. The variables in the equation are substituted if the model variable is input. + + Parameters + ---------- + equations: dict(~sympy.core.expr.Expr) + Dictionary of symbolic model equations. + params: QgParams + The parameters fully specifying the model configuration. + save_loc: str + Location to save the outputs as a .txt file. + language: str + Language syntax that the equations are returned in. Options are: + - `python` + - `fortran` + - `julia` + - `auto` + - `mathematica` + Default to `python`. + print_equations: bool + If `True`, equations are printed by the function, if `False`, equation strings are returned by the function. + Defaults to `False` + + Returns + ------- + equation_dict: dict(~sympy.core.expr.Expr) + Dictionary of symbolic model equations, that have been substituted with numerical values. + + """ + equation_dict = dict() + + # Substitute variable symbols + vector_subs = dict() + if language == 'python': + for i in range(1, params.ndim+1): + vector_subs['U_'+str(i)] = Symbol('U['+str(i-1)+']') + + if language == 'fortran' or language == 'auto': + for i in range(1, params.ndim+1): + vector_subs['U_'+str(i)] = Symbol('U('+str(i)+')') + + if language == 'julia': + for i in range(1, params.ndim+1): + vector_subs['U_'+str(i)] = Symbol('U['+str(i)+']') + + if language == 'mathematica': + for i in range(1, params.ndim+1): + vector_subs['U_'+str(i)] = Symbol('U('+str(i)+')') + + for k in equations.keys(): + if isinstance(equations[k], float): + eq = equations[k] + else: + eq = equations[k].subs(vector_subs) + eq = eq.evalf() + + if (language is not None) and print_equations: + eq = translate_equations(eq, language) + + equation_dict[k] = eq + + if print_equations: + if save_loc is None: + for eq in equation_dict.values(): + if save_loc is None: + print(eq) + else: + with open(save_loc, 'w') as f: + for eq in equation_dict.values(): + f.write("%s\n" % eq) + print("Equations written") + else: + return equation_dict + + +def equations_to_string(equations): + """ + Converts the symbolic equations, held in a dict format, to a dict of strings. + + Parameters + ---------- + equations: dict(~sympy.core.expr.Expr) + Dictionary of the substituted symbolic model equations. + + Returns + ------- + dict(~string) + Dictionary of the substituted symbolic model equations. + """ + + str_eq = dict() + for key in equations.keys(): + str_eq[key] = str(equations[key]) + return str_eq + + +def equation_as_function(equations, params, language='python', continuation_variables=None): + """Converts the symbolic equations to a function in string format in the language syntax specified, + or a lambdified python function. + + Parameters + ---------- + equations: dict(~sympy.core.expr.Expr) + Dictionary of the substituted symbolic model equations. + params: QgParams + The parameters fully specifying the model configuration. + language: str + Language syntax that the equations are returned in. Options are: + - `python` + - `fortran` + - `julia` + - `auto` + - `mathematica` + Default to `python`. + continuation_variables: Iterable(Parameter, ScalingParameter, ParametersArray) + Variables that are not substituted with numerical values. If `None`, all symbols are substituted. + + Returns + ------- + f_output: callable or str + Output is a function as a string in the specified language syntax. + + """ + if continuation_variables is None: + continuation_variables = list() + + eq_dict = format_equations(equations, params, language=language) + eq_dict = equations_to_string(eq_dict) + + f_output = list() + if language == 'python': + f_output.append('@njit') + func_def_str = 'def f(t, U' + for v in continuation_variables: + func_def_str += ', ' + str(v.symbol) + + f_output.append(func_def_str + '):' ) + + f_output.append('\t# Tendency function of the qgs model') + for v in continuation_variables: + f_output.append('\t# ' + str(v.symbol) + ":\t" + str(v.description)) + + f_output.append('') + f_output.append('\tF = np.empty_like(U)') + + for n, eq in eq_dict.items(): + f_output.append('\tF['+str(n-1)+'] = ' + eq) + + f_output.append('\treturn F') + f_output = translate_equations(f_output, language='python') + f_output = '\n'.join(f_output) + + if language == 'julia': + f_output.append('function f!(du, U, p, t)') + f_output.append('\t# Tendency function of the qgs model') + + for i, v in enumerate(continuation_variables): + f_output.append('\t' + str(v.symbol) + " = p[" + str(i+1) + "] " + "\t# " + str(v.description)) + + f_output.append('') + for n, eq in eq_dict.items(): + f_output.append('\tdu['+str(n)+'] = ' + eq) + + f_output.append('end') + f_output = translate_equations(f_output, language='julia') + f_output = '\n'.join(f_output) + + if language == 'fortran': + f_var = '' + for fv in continuation_variables: + f_var += ', ' + str(fv.symbol) + f_output.append('SUBROUTINE FUNC(NDIM, t, U, F' + f_var + ')') + + f_output.append('\t! Tendency function of the qgs model') + f_output.append('\tINTEGER, INTENT(IN) :: NDIM') + f_output.append('\tDOUBLE PRECISION, INTENT(IN) :: U(NDIM), PAR(*)') + f_output.append('\tDOUBLE PRECISION, INTENT(OUT) :: F(NDIM)') + + for v in continuation_variables: + f_output.append('\tDOUBLE PRECISION, INTENT(IN) :: ' + str(v.symbol) + "\t! " + str(v.description)) + + f_output.append('') + + f_output = _split_equations(eq_dict, f_output) + + f_output.append('END SUBROUTINE') + f_output = translate_equations(f_output, language='fortran') + f_output = '\n'.join(f_output) + + if language == 'auto': + eq_dict = _split_equations(eq_dict, f_output) + auto_file, auto_config = create_auto_file(eq_dict, params, continuation_variables) + auto_file, auto_config = ( + translate_equations(auto_file, language='fortran'), translate_equations(auto_config, language='fortran')) + f_output = ['\n'.join(auto_file), '\n'.join(auto_config)] + + if language == 'mathematica': + # TODO: This function needs testing before release + f_output.append('F = Array[' + str(len(eq_dict)) + ']') + + for n, eq in eq_dict.items(): + f_output.append('F['+str(n)+'] = ' + str(eq)) + + # TODO !!!! Killing output as I have not tested the above code !!!! + eq_dict = translate_equations(eq_dict, language='mathematica') + f_output = '\n'.join(f_output) + f_output = None + + return f_output + + +def jacobian_as_function(equations, params, language='python', continuation_variables=None): + """Converts the symbolic equations of the jacobain to a function in string format in the language syntax specified, + or a lambdified python function. + + Parameters + ---------- + equations: dict(~sympy.core.expr.Expr) + Dictionary of the substituted symbolic model equations. + params: QgParams + The parameters fully specifying the model configuration. + language: str + Language syntax that the equations are returned in. Options are: + - `python` + - `fortran` + - `julia` + - `auto` + - `mathematica` + Default to `python`. + continuation_variables: Iterable(Parameter, ScalingParameter, ParametersArray) + Variables that are not substituted with numerical values. If `None`, all symbols are substituted. + + Returns + ------- + f_output: callable or str + Output is a function as a string in the specified language syntax + + """ + if continuation_variables is None: + continuation_variables = list() + + eq_dict = format_equations(equations, params, language=language) + eq_dict = equations_to_string(eq_dict) + + f_output = list() + if language == 'python': + f_output.append('@njit') + func_def_str = 'def jac(t, U' + for v in continuation_variables: + func_def_str += ', ' + str(v.symbol) + + f_output.append(func_def_str + '):') + + f_output.append('\t# Jacobian function of the qgs model') + + for v in continuation_variables: + f_output.append('\t# ' + str(v.symbol) + ":\t" + str(v.description)) + + f_output.append('') + f_output.append('\tJ = np.zeros((len(U), len(U)))') + for n, eq in eq_dict.items(): + f_output.append('\tJ[' + str(n[0] - 1) + ', ' + str(n[1] - 1) + '] = ' + str(eq)) + + f_output.append('\treturn J') + f_output = '\n'.join(f_output) + + if language == 'julia': + f_output.append('function jac!(du, U, p, t)') + f_output.append('\t# Jacobian function of the qgs model') + + for i, v in enumerate(continuation_variables): + f_output.append('\t' + str(v.symbol) + " = p[" + str(i+1) + "]") + + f_output.append('') + for n, eq in eq_dict.items(): + f_output.append('\tdu[' + str(n[0]) + ', ' + str(n[1]) + '] = ' + str(eq)) + + f_output.append('end') + eq_dict = translate_equations(eq_dict, language='julia') + f_output = '\n'.join(f_output) + + if language == 'fortran': + f_var = '' + + for fv in continuation_variables: + f_var += ', ' + str(fv.symbol) + f_output.append('SUBROUTINE FUNC(NDIM, t, U, JAC' + f_var + ')') + + f_output.append('\t! Jacobian function of the qgs model') + f_output.append('\tINTEGER, INTENT(IN) :: NDIM') + f_output.append('\tDOUBLE PRECISION, INTENT(IN) :: U(NDIM), PAR(*)') + f_output.append('\tDOUBLE PRECISION, INTENT(OUT) :: JAC(NDIM, NDIM)') + + for v in continuation_variables: + f_output.append('\tDOUBLE PRECISION, INTENT(IN) :: ' + str(v.symbol) + "\t! " + str(v.description)) + + f_output.append('') + + f_output = _split_equations(eq_dict, f_output, two_dim=True) + + f_output.append('END SUBROUTINE') + eq_dict = translate_equations(eq_dict, language='fortran') + f_output = '\n'.join(f_output) + + if language == 'auto': + eq_dict = _split_equations(eq_dict, f_output, two_dim=True) + auto_file, auto_config = create_auto_file(eq_dict, params, continuation_variables) + auto_file, auto_config = ( + translate_equations(auto_file, language='fortran'), translate_equations(auto_config, language='fortran')) + f_output = ['\n'.join(auto_file), '\n'.join(auto_config)] + + if language == 'mathematica': + # TODO: This function needs testing before release + + f_output.append('jac = Array[' + str(len(eq_dict)) + ']') + + for n, eq in eq_dict.items(): + f_output.append('jac[' + str(n[0]) + ', ' + str(n[1]) + '] = ' + eq) + + # TODO !!!! Killing output as I have not tested the above code !!!! + eq_dict = translate_equations(eq_dict, language='mathematica') + f_output = '\n'.join(f_output) + f_output = None + + return f_output + + +def create_auto_file(equations, params, continuation_variables, auto_main_template=None, auto_c_template=None, + initialize_params=False, initialize_solution=False): + """Creates the auto configuration file and the model file. + Saves files to specified folder. + + Parameters + ---------- + equations: dict + Dictionary of the substituted symbolic model equations + params: QgParams + The parameters fully specifying the model configuration. + continuation_variables: Iterable(Parameter, ScalingParameter, ParametersArray) + Variables that are not substituted with numerical values. If `None`, all symbols are substituted + auto_main_template: str, optional + The template to be used to generate the main AUTO file. + If not provided, use the default template. + auto_c_template: str, optional + The template to be used to generate the AUTO config file. + If not provided, use the default template. + initialize_params: bool, optional + Add lines in the AUTO STPNT function to initialize the parameters. Default to `False`. + initialize_solution: bool, optional + Add lines in the AUTO STPNT function to initialize the solution. Default to `False`. + + Returns + ------- + auto_file: Str + The auto model file as a string + + auto_config: Str + Auto configuration file as a string + """ + + # TODO: There is some weird double tab spacings in the output, and I am not sure why + + # User passes the equations, with the variables to leave as variables. + # The existing model parameters are used to populate the auto file + # The variables given as `continuation_variables` remain in the equations. + # There is a limit of 1-10 remaining variables + + if (len(continuation_variables) < 1) or (len(continuation_variables) > 10): + raise ValueError("Too many variables for auto file") + + # Declare variables + declare_var = list() + for v in continuation_variables: + declare_var.append('DOUBLE PRECISION ' + str(v.symbol)) + + # make list of parameters + var_list = list() + var_ini = list() + sol_ini = list() + + for i, v in enumerate(continuation_variables): + temp_str = str(v.symbol) + " = PAR(" + str(i+1) + ")" + initial_value = "PAR(" + str(i+1) + ") = " + str(v) + " ! Variable: " + str(v.symbol) + + var_list.append(temp_str) + var_ini.append(initial_value) + + for i in range(1, params.ndim+1): + initial_sol = "U(" + str(i) + ") = 0.0d0" + + sol_ini.append(initial_sol) + + # Writing model file ################ + + if auto_main_template is not None: + lines = auto_main_template.split('\n') + else: + lines = default_auto_main_template.split('\n') + + auto_file = list() + for ln in lines: + if 'PARAMETER DECLARATION' in ln: + for dv in declare_var: + auto_file.append('\t' + dv) + elif 'CONTINUATION PARAMETERS' in ln: + for v in var_list: + auto_file.append('\t' + v) + elif 'EVOLUTION EQUATIONS' in ln: + for e in equations: + auto_file.append(e) + elif 'INITIALISE PARAMETERS' in ln and initialize_params: + for iv in var_ini: + auto_file.append('\t' + iv) + elif 'INITIALISE SOLUTION' in ln and initialize_solution: + for iv in sol_ini: + auto_file.append('\t' + iv) + else: + auto_file.append(ln.replace('\n', '')) + + # Writing config file ################ + + if auto_c_template is not None: + lines = auto_c_template.split('\n') + else: + lines = default_auto_c_template.split('\n') + + auto_config = list() + for ln in lines: + if '! PARAMETERS' in ln: + params_dic = {i+1: str(v.symbol) for i, v in enumerate(continuation_variables)} + params_dic.update({11: 'T', 12: 'theta', 14: 't', 25: 'T_r'}) + auto_config.append('parnames = ' + str(params_dic)) + + elif '! VARIABLES' in ln: + auto_config.append('unames = ' + str({i+1: params.var_string[i] for i in range(params.ndim)})) + + elif '! DIMENSION' in ln: + auto_config.append('NDIM = ' + str(params.ndim)) + + elif '! CONTINUATION ORDER' in ln: + auto_config.append('ICP = ' + str([str(v.symbol) for v in continuation_variables])) + + elif '! SOLUTION SAVE' in ln: + auto_config.append("# ! User to input save locations") + auto_config.append('UZR = ' + str({str(v.symbol): [] for v in continuation_variables})) + + elif '! STOP CONDITIONS' in ln: + auto_config.append("# ! User to input variable bounds") + auto_config.append('UZSTOP = ' + str({str(v.symbol): [] for v in continuation_variables})) + + else: + auto_config.append(ln.replace('\n', '')) + + return auto_file, auto_config + + +def _split_equations(eq_dict, f_output, line_len=80, two_dim=False): + """Function to split FORTRAN equations to a set length when producing functions""" + + for n, eq in eq_dict.items(): + # split equations to be a maximum of `line_len` + # split remainder of equation into chunks of length `line_length` + + # First translate the equation to ensure variable names are not split across rows + eq_translated = translate_equations(eq, language='fortran') + eq_chunks = [eq_translated[x: x + line_len] for x in range(0, len(eq_translated), line_len)] + if len(eq_chunks) > 1: + if two_dim: + f_output.append('\tJAC(' + str(n[0]) + ', ' + str(n[1]) + ') =\t ' + eq_chunks[0] + "&") + else: + f_output.append('\tF(' + str(n) + ') =\t ' + eq_chunks[0] + "&") + for ln in eq_chunks[1:-1]: + f_output.append("\t\t&" + ln + "&") + + f_output.append("\t\t&" + eq_chunks[-1]) + else: + if two_dim: + f_output.append('\tJAC(' + str(n[0]) + ', ' + str(n[1]) + ') =\t ' + eq_chunks[0]) + else: + f_output.append('\tF(' + str(n) + ') =\t ' + eq_chunks[0]) + f_output.append('') + return f_output + + +# ------------- Default AUTO files templates ---------------- + +default_auto_main_template = """!---------------------------------------------------------------------- +!---------------------------------------------------------------------- +! AUTO file for qgs model +!---------------------------------------------------------------------- +!---------------------------------------------------------------------- + +SUBROUTINE FUNC(NDIM,U,ICP,PAR,IJAC,F,DFDU,DFDP) +\t!--------- ---- + +\t! Evaluates the algebraic equations or ODE right hand side + +\t! Input arguments : +\t! NDIM : Dimension of the algebraic or ODE system +\t! U : State variables +\t! ICP : Array indicating the free parameter(s) +\t! PAR : Equation parameters + +\t! Values to be returned : +\t! F : Equation or ODE right hand side values + +\t! Normally unused Jacobian arguments : IJAC, DFDU, DFDP (see manual) + +\tIMPLICIT NONE +\tINTEGER, INTENT(IN) :: NDIM, IJAC, ICP(*) +\tDOUBLE PRECISION, INTENT(IN) :: U(NDIM), PAR(*) +\tDOUBLE PRECISION, INTENT(OUT) :: F(NDIM) +\tDOUBLE PRECISION, INTENT(INOUT) :: DFDU(NDIM,NDIM),DFDP(NDIM,*) + +! PARAMETER DECLARATION + +! CONTINUATION PARAMETERS + + +! EVOLUTION EQUATIONS + +END SUBROUTINE FUNC + +!----------------------------------------------------------------------- +!----------------------------------------------------------------------- + +SUBROUTINE STPNT(NDIM,U,PAR,T) +\t!--------- ----- + +\t! Input arguments : +\t! NDIM : Dimension of the algebraic or ODE system + +\t! Values to be returned : +\t! U : A starting solution vector +\t! PAR : The corresponding equation-parameter values + +\t! Note : For time- or space-dependent solutions this subroutine has +\t! the scalar input parameter T contains the varying time or space +\t! variable value. + +\tIMPLICIT NONE +\tINTEGER, INTENT(IN) :: NDIM +\tDOUBLE PRECISION, INTENT(INOUT) :: U(NDIM),PAR(*) +\tDOUBLE PRECISION, INTENT(IN) :: T +\tDOUBLE PRECISION :: X(NDIM+1) +\tINTEGER :: i,is + +\t! Initialize the equation parameters + +! INITIALISE PARAMETERS + +\t! Initialize the solution + +! INITIALISE SOLUTION + +\t! Initialization from a solution file (selection with PAR36) +\t! open(unit=15,file='',status='old') +\t! is=int(PAR(36)) +\t! if (is.gt.0) print*, 'Loading from solution :',is +\t! DO i=1,is +\t! read(15,*) X +\t! ENDDO +\t! close(15) +\t! U=X(2:NDIM+1) + +END SUBROUTINE STPNT + +!---------------------------------------------------------------------- +!---------------------------------------------------------------------- + +SUBROUTINE BCND(NDIM,PAR,ICP,NBC,U0,U1,FB,IJAC,DBC) +\t!--------- ---- + +\t! Boundary Conditions + +\t! Input arguments : +\t! NDIM : Dimension of the ODE system +\t! PAR : Equation parameters +\t! ICP : Array indicating the free parameter(s) +\t! NBC : Number of boundary conditions +\t! U0 : State variable values at the left boundary +\t! U1 : State variable values at the right boundary + +\t! Values to be returned : +\t! FB : The values of the boundary condition functions + +\t! Normally unused Jacobian arguments : IJAC, DBC (see manual) + +\tIMPLICIT NONE +\tINTEGER, INTENT(IN) :: NDIM, ICP(*), NBC, IJAC +\tDOUBLE PRECISION, INTENT(IN) :: PAR(*), U0(NDIM), U1(NDIM) +\tDOUBLE PRECISION, INTENT(OUT) :: FB(NBC) +\tDOUBLE PRECISION, INTENT(INOUT) :: DBC(NBC,*) + +\t!X FB(1)= +\t!X FB(2)= + +END SUBROUTINE BCND + +!---------------------------------------------------------------------- +!---------------------------------------------------------------------- + +SUBROUTINE ICND(NDIM,PAR,ICP,NINT,U,UOLD,UDOT,UPOLD,FI,IJAC,DINT) + +\t! Integral Conditions + +\t! Input arguments : +\t! NDIM : Dimension of the ODE system +\t! PAR : Equation parameters +\t! ICP : Array indicating the free parameter(s) +\t! NINT : Number of integral conditions +\t! U : Value of the vector function U at `time' t + +\t! The following input arguments, which are normally not needed, +\t! correspond to the preceding point on the solution branch +\t! UOLD : The state vector at 'time' t +\t! UDOT : Derivative of UOLD with respect to arclength +\t! UPOLD : Derivative of UOLD with respect to `time' + +\t! Normally unused Jacobian arguments : IJAC, DINT + +\t! Values to be returned : +\t! FI : The value of the vector integrand + +\tIMPLICIT NONE +\tINTEGER, INTENT(IN) :: NDIM, ICP(*), NINT, IJAC +\tDOUBLE PRECISION, INTENT(IN) :: PAR(*) +\tDOUBLE PRECISION, INTENT(IN) :: U(NDIM), UOLD(NDIM), UDOT(NDIM), UPOLD(NDIM) +\tDOUBLE PRECISION, INTENT(OUT) :: FI(NINT) +\tDOUBLE PRECISION, INTENT(INOUT) :: DINT(NINT,*) + +END SUBROUTINE ICND + +!---------------------------------------------------------------------- +!---------------------------------------------------------------------- + + +SUBROUTINE FOPT(NDIM,U,ICP,PAR,IJAC,FS,DFDU,DFDP) +\t!--------- ---- +\t! +\t! Defines the objective function for algebraic optimization problems +\t! +\t! Supplied variables : +\t! NDIM : Dimension of the state equation +\t! U : The state vector +\t! ICP : Indices of the control parameters +\t! PAR : The vector of control parameters +\t! +\t! Values to be returned : +\t! FS : The value of the objective function +\t! +\t! Normally unused Jacobian argument : IJAC, DFDP + +\tIMPLICIT NONE +\tINTEGER, INTENT(IN) :: NDIM, ICP(*), IJAC +\tDOUBLE PRECISION, INTENT(IN) :: U(NDIM), PAR(*) +\tDOUBLE PRECISION, INTENT(OUT) :: FS +\tDOUBLE PRECISION, INTENT(INOUT) :: DFDU(NDIM),DFDP(*) + +END SUBROUTINE FOPT + +!---------------------------------------------------------------------- +!---------------------------------------------------------------------- + +SUBROUTINE PVLS(NDIM,U,PAR) +\t!--------- ---- + +\tIMPLICIT NONE +\tINTEGER, INTENT(IN) :: NDIM +\tDOUBLE PRECISION, INTENT(INOUT) :: U(NDIM) +\tDOUBLE PRECISION, INTENT(INOUT) :: PAR(*) +\tDOUBLE PRECISION :: GETP,pi,realfm,imagfm,imagfm1 +\tDOUBLE PRECISION :: lw,lw1 +\tLOGICAL, SAVE :: first = .TRUE. +\tDOUBLE PRECISION :: T +\tINTEGER :: i + +\t!IF (first) THEN +\t\t!CALL STPNT(NDIM,U,PAR,T) +\t\t!first = .FALSE. +\t!ENDIF + +\tPAR(25)=0. +\tpi = 4*ATAN(1d0) +\ti=1 +\tlw=100. +\tlw1=101. +\tDO WHILE(i < NDIM) +\t\trealfm = GETP('EIG',I*2-1,U) +\t\tIF (ABS(realfm) < lw) THEN +\t\t\tlw = ABS(realfm) +\t\t\tlw1 = ABS(GETP('EIG',(I+1)*2-1,U)) +\t\t\timagfm1 = ABS(GETP('EIG',(I+1)*2,U)) +\t\t\timagfm = ABS(GETP('EIG',I*2,U)) +\t\tEND IF +\t\ti=i+1 +\tEND DO +\tIF ((lw==lw1).AND.(imagfm1==imagfm).AND.(imagfm/=0.D0)) THEN +\tPAR(25) = 2*pi/imagfm +\tENDIF +\t!---------------------------------------------------------------------- +\t! NOTE : +\t! Parameters set in this subroutine should be considered as ``solution +\t! measures'' and be used for output purposes only. +\t! +\t! They should never be used as `true'' continuation parameters. +\t! +\t! They may, however, be added as ``over-specified parameters'' in the +\t! parameter list associated with the AUTO-Constant NICP, in order to +\t! print their values on the screen and in the ``p.xxx file. +\t! +\t! They may also appear in the list associated with AUTO-Constant NUZR. +\t! +\t!---------------------------------------------------------------------- +\t! For algebraic problems the argument U is, as usual, the state vector. +\t! For differential equations the argument U represents the approximate +\t! solution on the entire interval [0,1]. In this case its values must +\t! be accessed indirectly by calls to GETP, as illustrated below. +\t!---------------------------------------------------------------------- +\t! +\t! Set PAR(2) equal to the L2-norm of U(1) +\t!X PAR(2)=GETP('NRM',1,U) +\t! +\t! Set PAR(3) equal to the minimum of U(2) +\t!X PAR(3)=GETP('MIN',2,U) +\t! +\t! Set PAR(4) equal to the value of U(2) at the left boundary. +\t!X PAR(4)=GETP('BV0',2,U) +\t! +\t! Set PAR(5) equal to the pseudo-arclength step size used. +\t!X PAR(5)=GETP('STP',1,U) +\t! +\t!---------------------------------------------------------------------- +\t! The first argument of GETP may be one of the following: +\t! 'NRM' (L2-norm), 'MAX' (maximum), +\t! 'INT' (integral), 'BV0 (left boundary value), +\t! 'MIN' (minimum), 'BV1' (right boundary value). +\t! +\t! Also available are +\t! 'STP' (Pseudo-arclength step size used). +\t! 'FLD' (`Fold function', which vanishes at folds). +\t! 'BIF' (`Bifurcation function', which vanishes at singular points). +\t! 'HBF' (`Hopf function'; which vanishes at Hopf points). +\t! 'SPB' ( Function which vanishes at secondary periodic bifurcations). +\t!---------------------------------------------------------------------- +END SUBROUTINE PVLS +""" + +default_auto_c_template = """#Configuration files + +#Parameters name +# ! PARAMETERS +#Variables name +# ! VARIABLES +#Dimension of the system +# ! DIMENSION +#Problem type (1 for FP, 2 for PO, -2 for time integration) +IPS = 1 +#Start solution label +IRS = 0 +#Continuation parameters (in order of use) +# ! CONTINUATION ORDER +#Number of mesh intervals +NTST= 100 +#Print and restart every NPR steps (0 to disable) +NPR= 0 +#Number of bifurcating branches to compute (negative number means continue only in one direction) +MXBF=0 +#Detection of Special Points +ISP=2 +#Maximum number of iteration in the Newton-Chord method +ITNW=7 +#Arc-length continuation parameters +DS = 0.00001, DSMIN= 1e-15, DSMAX= 1.0 +#Precision parameters (Typiq. EPSS = EPSL * 10^-2) +EPSL=1e-07, EPSU=1e-07, EPSS=1e-05 +#Number of parameter (don't change it) +NPAR = 36 +#User defined value where to save the solution +# ! SOLUTION SAVE +#Stop conditions +# ! STOP CONDITIONS +""" diff --git a/qgs/inner_products/analytic.py b/qgs/inner_products/analytic.py index 036c483..fa40f6b 100644 --- a/qgs/inner_products/analytic.py +++ b/qgs/inner_products/analytic.py @@ -41,7 +41,8 @@ from qgs.basis.fourier import channel_wavenumbers, basin_wavenumbers from qgs.inner_products.base import AtmosphericInnerProducts, OceanicInnerProducts, GroundInnerProducts -# TODO: Add warnings if trying to connect analytic and symbolic inner products together +# TODO: - Add warnings if trying to connect analytic and symbolic inner products together +# - Allow analytic inner product to be returned as symbolic arrays class AtmosphericAnalyticInnerProducts(AtmosphericInnerProducts): diff --git a/qgs/inner_products/definition.py b/qgs/inner_products/definition.py index 86187be..84ccce0 100644 --- a/qgs/inner_products/definition.py +++ b/qgs/inner_products/definition.py @@ -13,9 +13,9 @@ # from sympy.simplify import trigsimp from sympy.simplify.fu import TR8, TR10 -from sympy import diff, integrate, symbols, pi, Integral +from sympy import Symbol, diff, integrate, symbols, pi, Integral -_n = symbols('n', real=True, nonnegative=True) +_n = Symbol('n', positive=True) _x, _y = symbols('x y') @@ -95,7 +95,7 @@ def ip_lap(self, S, G, symbolic_expr=False): @abstractmethod def ip_diff_x(self, S, G, symbolic_expr=False): - """Function to compute the inner product :math:`(S, \partial_x G)`. + """Function to compute the inner product :math:`(S, \\partial_x G)`. Parameters ---------- @@ -215,7 +215,7 @@ def ip_lap(self, S, G, symbolic_expr=False, integrand=False): return self.symbolic_inner_product(S, self.laplacian(G), symbolic_expr=symbolic_expr, integrand=integrand) def ip_diff_x(self, S, G, symbolic_expr=False, integrand=False): - """Function to compute the inner product :math:`(S, \partial_x G)`. + """Function to compute the inner product :math:`(S, \\partial_x G)`. Parameters ---------- @@ -285,7 +285,7 @@ def ip_jac_lap(self, S, G, H, symbolic_expr=False, integrand=False): class StandardSymbolicInnerProductDefinition(SymbolicInnerProductDefinition): - """Standard `qgs` class to define symbolic inner products using `Sympy`_. + """Standard qgs class to define symbolic inner products using `Sympy`_. Parameters ---------- @@ -323,7 +323,7 @@ def jacobian(S, G): .. math: - J(S, G) = \partial_x S\, \partial_y G - \partial_y S\, \partial_x G + J(S, G) = \\partial_x S\\, \\partial_y G - \\partial_y S\\, \\partial_x G Parameters ---------- @@ -359,7 +359,7 @@ def laplacian(S): @staticmethod def integrate_over_domain(expr, symbolic_expr=False): """Definition of the normalized integrals over the spatial domain used by the inner products: - :math:`\\frac{n}{2\\pi^2}\\int_0^\\pi\\int_0^{2\\pi/n} \, \\mathrm{expr}(x, y) \, \\mathrm{d} x \, \\mathrm{d} y`. + :math:`\\frac{n}{2\\pi^2}\\int_0^\\pi\\int_0^{2\\pi/n} \\, \\mathrm{expr}(x, y) \\, \\mathrm{d} x \\, \\mathrm{d} y`. Parameters ---------- @@ -380,7 +380,7 @@ def integrate_over_domain(expr, symbolic_expr=False): def symbolic_inner_product(self, S, G, symbolic_expr=False, integrand=False): """Function defining the inner product to be computed symbolically: - :math:`(S, G) = \\frac{n}{2\\pi^2}\\int_0^\\pi\\int_0^{2\\pi/n} S(x,y)\, G(x,y)\, \\mathrm{d} x \, \\mathrm{d} y`. + :math:`(S, G) = \\frac{n}{2\\pi^2}\\int_0^\\pi\\int_0^{2\\pi/n} S(x,y)\\, G(x,y)\\, \\mathrm{d} x \\, \\mathrm{d} y`. Parameters ---------- @@ -402,4 +402,4 @@ def symbolic_inner_product(self, S, G, symbolic_expr=False, integrand=False): if integrand: return expr, (_x, 0, 2 * pi / _n), (_y, 0, pi) else: - return self.integrate_over_domain(self.optimizer(expr), symbolic_expr=symbolic_expr) \ No newline at end of file + return self.integrate_over_domain(self.optimizer(expr), symbolic_expr=symbolic_expr) diff --git a/qgs/inner_products/symbolic.py b/qgs/inner_products/symbolic.py index 323d269..ccc6dd9 100644 --- a/qgs/inner_products/symbolic.py +++ b/qgs/inner_products/symbolic.py @@ -31,14 +31,11 @@ from qgs.params.params import QgParams from qgs.inner_products.base import AtmosphericInnerProducts, OceanicInnerProducts, GroundInnerProducts from qgs.inner_products.definition import StandardSymbolicInnerProductDefinition -from sympy import lambdify from scipy.integrate import dblquad -from sympy import symbols +from sympy import ImmutableSparseMatrix, ImmutableSparseNDimArray, lambdify # TODO: - Add warnings if trying to connect analytic and symbolic inner products together -_n = symbols('n', real=True, nonnegative=True) - class AtmosphericSymbolicInnerProducts(AtmosphericInnerProducts): """Class which contains all the atmospheric inner products coefficients needed for the tendencies @@ -50,7 +47,7 @@ class AtmosphericSymbolicInnerProducts(AtmosphericInnerProducts): An instance of model's parameters object or a list in the form [aspect_ratio, atmospheric_basis, basis, oog, oro_basis]. If a list is provided, `aspect_ratio` is the aspect ratio of the domain, `atmospheric_basis` is a SymbolicBasis with the modes of the atmosphere, and `ocean_basis` is either `None` or a SymbolicBasis object with the modes of - the ocean or the ground. Finally `oog` indicates if it is an ocean or a ground component that is connected, + the ocean or the ground. Finally, `oog` indicates if it is an ocean or a ground component that is connected, by setting it to `ocean` or to 'ground', and in this latter case, `oro_basis` indicates on which basis the orography is developed. stored: bool, optional Indicate if the inner product must be stored or computed on the fly. Default to `True` @@ -104,7 +101,7 @@ class AtmosphericSymbolicInnerProducts(AtmosphericInnerProducts): """ def __init__(self, params=None, stored=True, inner_product_definition=None, interaction_inner_product_definition=None, - num_threads=None, quadrature=True, timeout=None, dynTinnerproducts=None, T4innerproducts=None): + num_threads=None, quadrature=True, timeout=None, dynTinnerproducts=None, T4innerproducts=None, return_symbolic=False, make_substitution=True): AtmosphericInnerProducts.__init__(self) @@ -173,7 +170,16 @@ def __init__(self, params=None, stored=True, inner_product_definition=None, inte self.ground_basis = None self.connected_to_ground = False - self.subs = [(_n, self.n)] + self.return_symbolic = return_symbolic + if return_symbolic: + self.mk_subs = make_substitution + if self.mk_subs: + self.subs = [(self.n.symbol, self.n)] + else: + self.subs = None + else: + self.mk_subs = True + self.subs = [(self.n.symbol, self.n)] if inner_product_definition is None: self.ip = StandardSymbolicInnerProductDefinition() @@ -235,45 +241,62 @@ def connect_to_ocean(self, ocean_basis, num_threads=None, timeout=None): num_threads = cpu_count() with Pool(max_workers=num_threads) as pool: - - subs = self.subs + self.atmospheric_basis.substitutions + self.oceanic_basis.substitutions + if self.mk_subs: + subs = self.subs + self.atmospheric_basis.substitutions + self.oceanic_basis.substitutions + else: + subs = self.subs noc = len(ocean_basis) - self._gh = None - self._d = sp.zeros((self.natm, noc), dtype=float, format='dok') - self._s = sp.zeros((self.natm, noc), dtype=float, format='dok') - if self._T4 or self._dynamic_T: - self._v = sp.zeros((self.natm, noc, noc, noc, noc), dtype=float, format='dok') + if self.return_symbolic: + self._gh = None + self._d = None + self._s = None + self._v = None + else: + self._gh = None + self._d = sp.zeros((self.natm, noc), dtype=float, format='dok') + self._s = sp.zeros((self.natm, noc), dtype=float, format='dok') + if self._T4 or self._dynamic_T: + self._v = sp.zeros((self.natm, noc, noc, noc, noc), dtype=float, format='dok') # d inner products args_list = [[(i, j), self.iip.ip_lap, (self._F(i), self._phi(j))] for i in range(self.natm) for j in range(noc)] - _parallel_compute(pool, args_list, subs, self._d, timeout) + output = _parallel_compute(pool, args_list, subs, self._d, timeout, symbolic_int=not self.mk_subs) + if self.return_symbolic: + self._d = ImmutableSparseMatrix(self.natm, noc, output) + else: + self._d = self._d.to_coo() # s inner products args_list = [[(i, j), self.iip.symbolic_inner_product, (self._F(i), self._phi(j))] for i in range(self.natm) for j in range(noc)] - _parallel_compute(pool, args_list, subs, self._s, timeout) + output = _parallel_compute(pool, args_list, subs, self._s, timeout, symbolic_int=not self.mk_subs) + if self.return_symbolic: + self._s = ImmutableSparseMatrix(self.natm, noc, output) + else: + self._s = self._s.to_coo() if self._T4: # v inner products args_list = [[(i, j, k, ell, m), self.ip.symbolic_inner_product, (self._F(i), self._phi(j) * self._phi(k) * self._phi(ell) * self._phi(m))] for i in range(self.natm) for j in range(noc) for k in range(j, noc) for ell in range(k, noc) for m in range(ell, noc)] - _parallel_compute(pool, args_list, subs, self._v, timeout, permute=True) + output = _parallel_compute(pool, args_list, subs, self._v, timeout, permute=True, symbolic_int=not self.mk_subs) elif self._dynamic_T: # v inner products args_list = [[(i, 0, 0, 0, m), self.ip.symbolic_inner_product, (self._F(i), self._phi(0) * self._phi(0) * self._phi(0) * self._phi(m))] for i in range(self.natm) for m in range(noc)] - _parallel_compute(pool, args_list, subs, self._v, timeout, permute=True) + output = _parallel_compute(pool, args_list, subs, self._v, timeout, permute=True, symbolic_int=not self.mk_subs) - self._s = self._s.to_coo() - self._d = self._d.to_coo() - if self._T4 or self._dynamic_T: - self._v = self._v.to_coo() + if self._T4 or self._dynamic_T: + if self.return_symbolic: + self._v = ImmutableSparseNDimArray(output, shape=(self.natm, noc, noc, noc, noc)) + else: + self._v = self._v.to_coo() def connect_to_ground(self, ground_basis, orographic_basis, num_threads=None, timeout=None): """Connect the atmosphere to the ground. @@ -309,50 +332,67 @@ def connect_to_ground(self, ground_basis, orographic_basis, num_threads=None, ti num_threads = cpu_count() with Pool(max_workers=num_threads) as pool: - - subs = self.subs + self.atmospheric_basis.substitutions + self.ground_basis.substitutions + if self.mk_subs: + subs = self.subs + self.atmospheric_basis.substitutions + self.ground_basis.substitutions + else: + subs = self.subs ngr = len(ground_basis) - if orographic_basis == "atmospheric": + if self.return_symbolic: self._gh = None + self._d = None + self._s = None + self._v = None else: - self._gh = sp.zeros((self.natm, self.natm, ngr), dtype=float, format='dok') - self._d = None - self._s = sp.zeros((self.natm, ngr), dtype=float, format='dok') - if self._T4 or self._dynamic_T: - self._v = sp.zeros((self.natm, ngr, ngr, ngr, ngr), dtype=float, format='dok') + if orographic_basis == "atmospheric": + self._gh = None + else: + self._gh = sp.zeros((self.natm, self.natm, ngr), dtype=float, format='dok') + self._d = None + self._s = sp.zeros((self.natm, ngr), dtype=float, format='dok') + if self._T4 or self._dynamic_T: + self._v = sp.zeros((self.natm, ngr, ngr, ngr, ngr), dtype=float, format='dok') # s inner products args_list = [[(i, j), self.iip.symbolic_inner_product, (self._F(i), self._phi(j))] for i in range(self.natm) for j in range(ngr)] - _parallel_compute(pool, args_list, subs, self._s, timeout) + output = _parallel_compute(pool, args_list, subs, self._s, timeout, symbolic_int=not self.mk_subs) + if self.return_symbolic: + self._s = ImmutableSparseMatrix(self.natm, ngr, output) + else: + self._s = self._s.to_coo() # gh inner products if orographic_basis != "atmospheric": args_list = [[(i, j, k), self.iip.ip_jac, (self._F(i), self._F(j), self._phi(k))] for i in range(self.natm) for j in range(self.natm) for k in range(ngr)] - _parallel_compute(pool, args_list, subs, self._gh, timeout) + output = _parallel_compute(pool, args_list, subs, self._gh, timeout, symbolic_int=not self.mk_subs) + if self.return_symbolic: + self._gh = ImmutableSparseNDimArray(output, shape=(self.natm, self.natm, ngr)) + else: + if self._gh is not None: + self._gh = self._gh.to_coo() if self._T4: # v inner products args_list = [[(i, j, k, ell, m), self.ip.symbolic_inner_product, (self._F(i), self._phi(j) * self._phi(k) * self._phi(ell) * self._phi(m))] for i in range(self.natm) for j in range(ngr) for k in range(j, ngr) for ell in range(k, ngr) for m in range(ell, ngr)] - _parallel_compute(pool, args_list, subs, self._v, timeout, permute=True) + output = _parallel_compute(pool, args_list, subs, self._v, timeout, permute=True, symbolic_int=not self.mk_subs) elif self._dynamic_T: # v inner products args_list = [[(i, 0, 0, 0, m), self.ip.symbolic_inner_product, (self._F(i), self._phi(0) * self._phi(0) * self._phi(0) * self._phi(m))] for i in range(self.natm) for m in range(ngr)] - _parallel_compute(pool, args_list, subs, self._v, timeout, permute=True) + output = _parallel_compute(pool, args_list, subs, self._v, timeout, permute=True, symbolic_int=not self.mk_subs) - self._s = self._s.to_coo() - if self._gh is not None: - self._gh = self._gh.to_coo() - if self._T4 or self._dynamic_T: - self._v = self._v.to_coo() + if self._T4 or self._dynamic_T: + if self.return_symbolic: + self._v = ImmutableSparseNDimArray(output, shape=(self.natm, ngr, ngr, ngr, ngr)) + else: + self._v = self._v.to_coo() def compute_inner_products(self, num_threads=None, timeout=None): """Function computing and storing all the inner products at once. @@ -368,73 +408,100 @@ def compute_inner_products(self, num_threads=None, timeout=None): If `None` or `False`, no timeout occurs. Default to `None`. """ - - self._a = sp.zeros((self.natm, self.natm), dtype=float, format='dok') - self._u = sp.zeros((self.natm, self.natm), dtype=float, format='dok') - self._c = sp.zeros((self.natm, self.natm), dtype=float, format='dok') - self._b = sp.zeros((self.natm, self.natm, self.natm), dtype=float, format='dok') - self._g = sp.zeros((self.natm, self.natm, self.natm), dtype=float, format='dok') - if self._T4 or self._dynamic_T: - self._z = sp.zeros((self.natm, self.natm, self.natm, self.natm, self.natm), dtype=float, format='dok') + if self.return_symbolic: + self._a = None + self._u = None + self._c = None + self._b = None + self._g = None + self._z = None + else: + self._a = sp.zeros((self.natm, self.natm), dtype=float, format='dok') + self._u = sp.zeros((self.natm, self.natm), dtype=float, format='dok') + self._c = sp.zeros((self.natm, self.natm), dtype=float, format='dok') + self._b = sp.zeros((self.natm, self.natm, self.natm), dtype=float, format='dok') + self._g = sp.zeros((self.natm, self.natm, self.natm), dtype=float, format='dok') + if self._T4 or self._dynamic_T: + self._z = sp.zeros((self.natm, self.natm, self.natm, self.natm, self.natm), dtype=float, format='dok') if self.stored: if num_threads is None: num_threads = cpu_count() with Pool(max_workers=num_threads) as pool: - - subs = self.subs + self.atmospheric_basis.substitutions + if self.mk_subs: + subs = self.subs + self.atmospheric_basis.substitutions + else: + subs = self.subs # a inner products args_list = [[(i, j), self.ip.ip_lap, (self._F(i), self._F(j))] for i in range(self.natm) for j in range(self.natm)] - - _parallel_compute(pool, args_list, subs, self._a, timeout) + + output = _parallel_compute(pool, args_list, subs, self._a, timeout, symbolic_int=not self.mk_subs) + if self.return_symbolic: + self._a = ImmutableSparseMatrix(self.natm, self.natm, output) + else: + self._a = self._a.to_coo() # u inner products args_list = [[(i, j), self.ip.symbolic_inner_product, (self._F(i), self._F(j))] for i in range(self.natm) for j in range(self.natm)] - _parallel_compute(pool, args_list, subs, self._u, timeout) + output = _parallel_compute(pool, args_list, subs, self._u, timeout, symbolic_int=not self.mk_subs) + if self.return_symbolic: + self._u = ImmutableSparseMatrix(self.natm, self.natm, output) + else: + self._u = self._u.to_coo() # c inner products args_list = [[(i, j), self.ip.ip_diff_x, (self._F(i), self._F(j))] for i in range(self.natm) for j in range(self.natm)] - _parallel_compute(pool, args_list, subs, self._c, timeout) + output = _parallel_compute(pool, args_list, subs, self._c, timeout, symbolic_int=not self.mk_subs) + if self.return_symbolic: + self._c = ImmutableSparseMatrix(self.natm, self.natm, output) + else: + self._c = self._c.to_coo() # b inner products args_list = [[(i, j, k), self.ip.ip_jac_lap, (self._F(i), self._F(j), self._F(k))] for i in range(self.natm) for j in range(self.natm) for k in range(self.natm)] - _parallel_compute(pool, args_list, subs, self._b, timeout) + output = _parallel_compute(pool, args_list, subs, self._b, timeout, symbolic_int=not self.mk_subs) + if self.return_symbolic: + self._b = ImmutableSparseNDimArray(output, shape=(self.natm, self.natm, self.natm)) + else: + self._b = self._b.to_coo() # g inner products args_list = [[(i, j, k), self.ip.ip_jac, (self._F(i), self._F(j), self._F(k))] for i in range(self.natm) for j in range(self.natm) for k in range(self.natm)] - _parallel_compute(pool, args_list, subs, self._g, timeout) + output = _parallel_compute(pool, args_list, subs, self._g, timeout, symbolic_int=not self.mk_subs) + if self.return_symbolic: + self._g = ImmutableSparseNDimArray(output, shape=(self.natm, self.natm, self.natm)) + else: + self._g = self._g.to_coo() if self._T4: # z inner products args_list = [[(i, j, k, ell, m), self.ip.symbolic_inner_product, (self._F(i), self._F(j) * self._F(k) * self._F(ell) * self._F(m))] for i in range(self.natm) for j in range(self.natm) for k in range(j, self.natm) for ell in range(k, self.natm) for m in range(ell, self.natm)] - _parallel_compute(pool, args_list, subs, self._z, timeout, permute=True) + output = _parallel_compute(pool, args_list, subs, self._z, timeout, permute=True, symbolic_int=not self.mk_subs) elif self._dynamic_T: # z inner products args_list = [[(i, 0, 0, 0, m), self.ip.symbolic_inner_product, (self._F(i), self._F(0) * self._F(0) * self._F(0) * self._F(m))] for i in range(self.natm) for m in range(self.natm)] - _parallel_compute(pool, args_list, subs, self._z, timeout, permute=True) + output = _parallel_compute(pool, args_list, subs, self._z, timeout, permute=True, symbolic_int=not self.mk_subs) - self._a = self._a.to_coo() - self._u = self._u.to_coo() - self._c = self._c.to_coo() - self._g = self._g.to_coo() - self._b = self._b.to_coo() - if self._T4 or self._dynamic_T: - self._z = self._z.to_coo() + if self._T4 or self._dynamic_T: + if self.return_symbolic: + self._z = ImmutableSparseNDimArray(output, shape=(self.natm, self.natm, self.natm, self.natm, self.natm)) + else: + self._z = self._z.to_coo() @property def natm(self): @@ -446,17 +513,28 @@ def natm(self): # !-----------------------------------------------------! def _integrate(self, subs, args): + if not self.mk_subs: + res = _apply(args) + return res[1] + if self.quadrature: res = _num_apply(args) return res[1] else: res = _apply(args)[1] - return float(res.subs(subs)) + if subs is not None: + return float(res.subs(subs)) + else: + return res def a(self, i, j): """Function to compute the matrix of the eigenvalues of the Laplacian (atmospheric): :math:`a_{i, j} = (F_i, {\\nabla}^2 F_j)`.""" if not self.stored: - subs = self.subs + self.atmospheric_basis.substitutions + if self.mk_subs: + subs = self.subs + self.atmospheric_basis.substitutions + else: + subs = self.subs + args = ((i, j), self.ip.ip_lap, (self._F(i), self._F(j)), subs) return self._integrate(subs, args) else: @@ -465,7 +543,11 @@ def a(self, i, j): def u(self, i, j): """Function to compute the matrix of inner product: :math:`u_{i, j} = (F_i, F_j)`.""" if not self.stored: - subs = self.subs + self.atmospheric_basis.substitutions + if self.mk_subs: + subs = self.subs + self.atmospheric_basis.substitutions + else: + subs = self.subs + args = ((i, j), self.ip.symbolic_inner_product, (self._F(i), self._F(j)), subs) return self._integrate(subs, args) else: @@ -474,7 +556,11 @@ def u(self, i, j): def b(self, i, j, k): """Function to compute the tensors holding the Jacobian inner products: :math:`b_{i, j, k} = (F_i, J(F_j, \\nabla^2 F_k))`.""" if not self.stored: - subs = self.subs + self.atmospheric_basis.substitutions + if self.mk_subs: + subs = self.subs + self.atmospheric_basis.substitutions + else: + subs = self.subs + args = ((i, j, k), self.ip.ip_jac_lap, (self._F(i), self._F(j), self._F(k)), subs) return self._integrate(subs, args) else: @@ -483,7 +569,11 @@ def b(self, i, j, k): def c(self, i, j): """Function to compute the matrix of beta terms for the atmosphere: :math:`c_{i,j} = (F_i, \\partial_x F_j)`.""" if not self.stored: - subs = self.subs + self.atmospheric_basis.substitutions + if self.mk_subs: + subs = self.subs + self.atmospheric_basis.substitutions + else: + subs = self.subs + args = ((i, j), self.ip.ip_diff_x, (self._F(i), self._F(j)), subs) return self._integrate(subs, args) else: @@ -492,7 +582,11 @@ def c(self, i, j): def g(self, i, j, k): """Function to compute tensors holding the Jacobian inner products: :math:`g_{i,j,k} = (F_i, J(F_j, F_k))`.""" if not self.stored: - subs = self.subs + self.atmospheric_basis.substitutions + if self.mk_subs: + subs = self.subs + self.atmospheric_basis.substitutions + else: + subs = self.subs + args = ((i, j, k), self.ip.ip_jac, (self._F(i), self._F(j), self._F(k)), subs) return self._integrate(subs, args) else: @@ -505,14 +599,17 @@ def gh(self, i, j, k): if self.stored and self._gh is not None: return self._gh[i, j, k] else: - if self.connected_to_ocean: - extra_subs = self.oceanic_basis.substitutions - elif self.connected_to_ground: - extra_subs = self.ground_basis.substitutions + if self.mk_subs: + if self.connected_to_ocean: + extra_subs = self.oceanic_basis.substitutions + elif self.connected_to_ground: + extra_subs = self.ground_basis.substitutions + else: + extra_subs = None + subs = self.subs + self.atmospheric_basis.substitutions + extra_subs else: - extra_subs = None + subs = self.subs - subs = self.subs + self.atmospheric_basis.substitutions + extra_subs args = ((i, j, k), self.iip.ip_jac, (self._F(i), self._F(j), self._phi(k)), subs) return self._integrate(subs, args) else: @@ -524,14 +621,17 @@ def s(self, i, j): if self.stored and self._s is not None: return self._s[i, j] else: - if self.connected_to_ocean: - extra_subs = self.oceanic_basis.substitutions - elif self.connected_to_ground: - extra_subs = self.ground_basis.substitutions + if self.mk_subs: + if self.connected_to_ocean: + extra_subs = self.oceanic_basis.substitutions + elif self.connected_to_ground: + extra_subs = self.ground_basis.substitutions + else: + extra_subs = None + subs = self.subs + self.atmospheric_basis.substitutions + extra_subs else: - extra_subs = None + subs = self.subs - subs = self.subs + self.atmospheric_basis.substitutions + extra_subs args = ((i, j), self.iip.symbolic_inner_product, (self._F(i), self._phi(j)), subs) return self._integrate(subs, args) else: @@ -543,14 +643,18 @@ def d(self, i, j): if self.stored and self._d is not None: return self._d[i, j] else: - if self.connected_to_ocean: - extra_subs = self.oceanic_basis.substitutions - elif self.connected_to_ground: - extra_subs = self.ground_basis.substitutions + if self.mk_subs: + if self.connected_to_ocean: + extra_subs = self.oceanic_basis.substitutions + elif self.connected_to_ground: + extra_subs = self.ground_basis.substitutions + else: + extra_subs = None + + subs = self.subs + self.atmospheric_basis.substitutions + extra_subs else: - extra_subs = None + subs = self.subs - subs = self.subs + self.atmospheric_basis.substitutions + extra_subs args = ((i, j), self.iip.ip_lap, (self._F(i), self._phi(j)), subs) return self._integrate(subs, args) else: @@ -559,7 +663,11 @@ def d(self, i, j): def z(self, i, j, k, l, m): """Function to compute the :math:`T^4` temperature forcing for the radiation lost by atmosphere to space & ground/ocean: :math:`z_{i,j,k,l,m} = (F_i, F_j F_k F_l F_m)`.""" if not self.stored: - subs = self.subs + self.atmospheric_basis.substitutions + if self.mk_subs: + subs = self.subs + self.atmospheric_basis.substitutions + else: + subs = self.subs + args = ((i, j, k, l, m), self.ip.symbolic_inner_product, (self._F(i), self._F(j) * self._F(k) * self._F(l) * self._F(m)), subs) if self.quadrature: res = _num_apply(args) @@ -573,7 +681,11 @@ def z(self, i, j, k, l, m): def v(self, i, j, k, l, m): """Function to compute the :math:`T^4` temperature forcing of the ocean on the atmosphere: :math:`v_{i,j,k,l,m} = (F_i, \\phi_j \\phi_k \\phi_l \\phi_m)`.""" if not self.stored: - subs = self.subs + self.atmospheric_basis.substitutions + if self.mk_subs: + subs = self.subs + self.atmospheric_basis.substitutions + else: + subs = self.subs + args = ((i, j, k, l, m), self.ip.symbolic_inner_product, (self._F(i), self._phi(j) * self._phi(k) * self._phi(l) * self._phi(m)), subs) if self.quadrature: res = _num_apply(args) @@ -647,7 +759,7 @@ class OceanicSymbolicInnerProducts(OceanicInnerProducts): symbolic computation. """ def __init__(self, params=None, stored=True, inner_product_definition=None, interaction_inner_product_definition=None, - num_threads=None, quadrature=True, timeout=None, dynTinnerproducts=None, T4innerproducts=None): + num_threads=None, quadrature=True, timeout=None, dynTinnerproducts=None, T4innerproducts=None, return_symbolic=False, make_substitution=True): OceanicInnerProducts.__init__(self) @@ -697,7 +809,16 @@ def __init__(self, params=None, stored=True, inner_product_definition=None, inte self.atmospheric_basis = None self.connected_to_atmosphere = False - self.subs = [(_n, self.n)] + self.return_symbolic = return_symbolic + if return_symbolic: + self.mk_subs = make_substitution + if self.mk_subs: + self.subs = [(self.n.symbol, self.n)] + else: + self.subs = None + else: + self.mk_subs = True + self.subs = [(self.n.symbol, self.n)] if inner_product_definition is None: self.ip = StandardSymbolicInnerProductDefinition() @@ -751,44 +872,62 @@ def connect_to_atmosphere(self, atmosphere_basis, num_threads=None, timeout=None num_threads = cpu_count() with Pool(max_workers=num_threads) as pool: - - subs = self.subs + self.oceanic_basis.substitutions + self.atmospheric_basis.substitutions + if self.mk_subs: + subs = self.subs + self.oceanic_basis.substitutions + self.atmospheric_basis.substitutions + else: + subs = self.subs natm = len(atmosphere_basis) - self._K = sp.zeros((self.noc, natm), dtype=float, format='dok') - self._W = sp.zeros((self.noc, natm), dtype=float, format='dok') - if self._T4 or self._dynamic_T: - self._Z = sp.zeros((self.noc, natm, natm, natm, natm), dtype=float, format='dok') + if self.return_symbolic: + self._K = None + self._W = None + self._Z = None + else: + self._K = sp.zeros((self.noc, natm), dtype=float, format='dok') + self._W = sp.zeros((self.noc, natm), dtype=float, format='dok') + if self._T4 or self._dynamic_T: + self._Z = sp.zeros((self.noc, natm, natm, natm, natm), dtype=float, format='dok') # K inner products args_list = [[(i, j), self.iip.ip_lap, (self._phi(i), self._F(j))] for i in range(self.noc) - for j in range(natm)] + for j in range(natm)] + + output = _parallel_compute(pool, args_list, subs, self._K, timeout, symbolic_int=not self.mk_subs) - _parallel_compute(pool, args_list, subs, self._K, timeout) + if self.return_symbolic: + self._K = ImmutableSparseMatrix(self.noc, natm, output) + else: + self._K = self._K.to_coo() # W inner products args_list = [[(i, j), self.iip.symbolic_inner_product, (self._phi(i), self._F(j))] for i in range(self.noc) - for j in range(natm)] + for j in range(natm)] - _parallel_compute(pool, args_list, subs, self._W, timeout) + output = _parallel_compute(pool, args_list, subs, self._W, timeout, symbolic_int=not self.mk_subs) + + if self.return_symbolic: + self._W = ImmutableSparseMatrix(self.noc, natm, output) + else: + self._W = self._W.to_coo() if self._T4: # Z inner products args_list = [[(i, j, k, ell, m), self.ip.symbolic_inner_product, (self._phi(i), self._F(j) * self._F(k) * self._F(ell) * self._F(m))] for i in range(self.noc) for j in range(natm) for k in range(j, natm) for ell in range(k, natm) for m in range(ell, natm)] - _parallel_compute(pool, args_list, subs, self._Z, timeout, permute=True) + output = _parallel_compute(pool, args_list, subs, self._Z, timeout, permute=True, symbolic_int=not self.mk_subs) elif self._dynamic_T: # Z inner products args_list = [[(i, 0, 0, 0, m), self.ip.symbolic_inner_product, (self._phi(i), self._F(0) * self._F(0) * self._F(0) * self._F(m))] for i in range(self.noc) for m in range(natm)] - _parallel_compute(pool, args_list, subs, self._Z, timeout, permute=True) + output = _parallel_compute(pool, args_list, subs, self._Z, timeout, permute=True, symbolic_int=not self.mk_subs) - self._K = self._K.to_coo() - self._W = self._W.to_coo() if self._T4 or self._dynamic_T: - self._Z = self._Z.to_coo() + if self.return_symbolic: + self._Z = ImmutableSparseNDimArray(output, shape=(self.noc, natm, natm, natm, natm)) + else: + self._Z = self._Z.to_coo() def compute_inner_products(self, num_threads=None, timeout=None): """Function computing and storing all the inner products at once. @@ -804,70 +943,97 @@ def compute_inner_products(self, num_threads=None, timeout=None): If `None` or `False`, no timeout occurs. Default to `None`. """ - - self._M = sp.zeros((self.noc, self.noc), dtype=float, format='dok') - self._U = sp.zeros((self.noc, self.noc), dtype=float, format='dok') - self._N = sp.zeros((self.noc, self.noc), dtype=float, format='dok') - self._O = sp.zeros((self.noc, self.noc, self.noc), dtype=float, format='dok') - self._C = sp.zeros((self.noc, self.noc, self.noc), dtype=float, format='dok') - if self._T4 or self._dynamic_T: - self._V = sp.zeros((self.noc, self.noc, self.noc, self.noc, self.noc), dtype=float, format='dok') + if self.return_symbolic: + self._M = None + self._U = None + self._N = None + self._O = None + self._C = None + self._V = None + else: + self._M = sp.zeros((self.noc, self.noc), dtype=float, format='dok') + self._U = sp.zeros((self.noc, self.noc), dtype=float, format='dok') + self._N = sp.zeros((self.noc, self.noc), dtype=float, format='dok') + self._O = sp.zeros((self.noc, self.noc, self.noc), dtype=float, format='dok') + self._C = sp.zeros((self.noc, self.noc, self.noc), dtype=float, format='dok') + if self._T4 or self._dynamic_T: + self._V = sp.zeros((self.noc, self.noc, self.noc, self.noc, self.noc), dtype=float, format='dok') if self.stored: if num_threads is None: num_threads = cpu_count() with Pool(max_workers=num_threads) as pool: - - subs = self.subs + self.oceanic_basis.substitutions + if self.mk_subs: + subs = self.subs + self.oceanic_basis.substitutions + else: + subs = self.subs # N inner products args_list = [[(i, j), self.ip.ip_diff_x, (self._phi(i), self._phi(j))] for i in range(self.noc) for j in range(self.noc)] - _parallel_compute(pool, args_list, subs, self._N, timeout) + output = _parallel_compute(pool, args_list, subs, self._N, timeout, symbolic_int=not self.mk_subs) + if self.return_symbolic: + self._N = ImmutableSparseMatrix(self.noc, self.noc, output) + else: + self._N = self._N.to_coo() # M inner products args_list = [[(i, j), self.ip.ip_lap, (self._phi(i), self._phi(j))] for i in range(self.noc) for j in range(self.noc)] - _parallel_compute(pool, args_list, subs, self._M, timeout) + output = _parallel_compute(pool, args_list, subs, self._M, timeout, symbolic_int=not self.mk_subs) + if self.return_symbolic: + self._M = ImmutableSparseMatrix(self.noc, self.noc, output) + else: + self._M = self._M.to_coo() # U inner products args_list = [[(i, j), self.ip.symbolic_inner_product, (self._phi(i), self._phi(j))] for i in range(self.noc) for j in range(self.noc)] - _parallel_compute(pool, args_list, subs, self._U, timeout) + output = _parallel_compute(pool, args_list, subs, self._U, timeout, symbolic_int=not self.mk_subs) + if self.return_symbolic: + self._U = ImmutableSparseMatrix(self.noc, self.noc, output) + else: + self._U = self._U.to_coo() # O inner products args_list = [[(i, j, k), self.ip.ip_jac, (self._phi(i), self._phi(j), self._phi(k))] for i in range(self.noc) for j in range(self.noc) for k in range(self.noc)] - _parallel_compute(pool, args_list, subs, self._O, timeout) + output = _parallel_compute(pool, args_list, subs, self._O, timeout, symbolic_int=not self.mk_subs) + if self.return_symbolic: + self._O = ImmutableSparseNDimArray(output, shape=(self.noc, self.noc, self.noc)) + else: + self._O = self._O.to_coo() # C inner products args_list = [[(i, j, k), self.ip.ip_jac_lap, (self._phi(i), self._phi(j), self._phi(k))] for i in range(self.noc) for j in range(self.noc) for k in range(self.noc)] - _parallel_compute(pool, args_list, subs, self._C, timeout) + output = _parallel_compute(pool, args_list, subs, self._C, timeout, symbolic_int=not self.mk_subs) + if self.return_symbolic: + self._C = ImmutableSparseNDimArray(output, shape=(self.noc, self.noc, self.noc)) + else: + self._C = self._C.to_coo() if self._T4: # V inner products args_list = [[(i, j, k, ell, m), self.ip.symbolic_inner_product, (self._phi(i), self._phi(j) * self._phi(k) * self._phi(ell) * self._phi(m))] for i in range(self.noc) for j in range(self.noc) for k in range(j, self.noc) for ell in range(k, self.noc) for m in range(ell, self.noc)] - _parallel_compute(pool, args_list, subs, self._V, timeout, permute=True) + output = _parallel_compute(pool, args_list, subs, self._V, timeout, permute=True, symbolic_int=not self.mk_subs) elif self._dynamic_T: # V inner products args_list = [[(i, 0, 0, 0, m), self.ip.symbolic_inner_product, (self._phi(i), self._phi(0) * self._phi(0) * self._phi(0) * self._phi(m))] for i in range(self.noc) for m in range(self.noc)] - _parallel_compute(pool, args_list, subs, self._V, timeout, permute=True) + output = _parallel_compute(pool, args_list, subs, self._V, timeout, permute=True, symbolic_int=not self.mk_subs) - self._M = self._M.to_coo() - self._U = self._U.to_coo() - self._N = self._N.to_coo() - self._O = self._O.to_coo() - self._C = self._C.to_coo() - if self._T4 or self._dynamic_T: - self._V = self._V.to_coo() + if self._T4 or self._dynamic_T: + if self.return_symbolic: + self._V = ImmutableSparseNDimArray(output, shape=(self.noc, self.noc, self.noc, self.noc, self.noc)) + else: + self._V = self._V.to_coo() @property def noc(self): @@ -879,6 +1045,10 @@ def noc(self): # !-----------------------------------------------------! def _integrate(self, subs, args): + if self.return_symbolic: + res = _apply(args) + return res[1] + if self.quadrature: res = _num_apply(args) return res[1] @@ -889,7 +1059,11 @@ def _integrate(self, subs, args): def M(self, i, j): """Function to compute the forcing of the ocean fields on the ocean: :math:`M_{i,j} = (\\phi_i, \\nabla^2 \\phi_j)`.""" if not self.stored: - subs = self.subs + self.oceanic_basis.substitutions + if self.mk_subs: + subs = self.subs + self.oceanic_basis.substitutions + else: + subs = self.subs + args = ((i, j), self.ip.ip_lap, (self._phi(i), self._phi(j)), subs) return self._integrate(subs, args) else: @@ -898,7 +1072,11 @@ def M(self, i, j): def U(self, i, j): """Function to compute the inner products: :math:`U_{i,j} = (\\phi_i, \\phi_j)`.""" if not self.stored: - subs = self.subs + self.oceanic_basis.substitutions + if self.mk_subs: + subs = self.subs + self.oceanic_basis.substitutions + else: + subs = self.subs + args = ((i, j), self.ip.symbolic_inner_product, (self._phi(i), self._phi(j)), subs) return self._integrate(subs, args) else: @@ -907,7 +1085,11 @@ def U(self, i, j): def N(self, i, j): """Function computing the beta term for the ocean: :math:`N_{i,j} = (\\phi_i, \\partial_x \\phi_j)`.""" if not self.stored: - subs = self.subs + self.oceanic_basis.substitutions + if self.mk_subs: + subs = self.subs + self.oceanic_basis.substitutions + else: + subs = self.subs + args = ((i, j), self.ip.ip_diff_x, (self._phi(i), self._phi(j)), subs) return self._integrate(subs, args) else: @@ -916,7 +1098,11 @@ def N(self, i, j): def O(self, i, j, k): """Function to compute the temperature advection term (passive scalar): :math:`O_{i,j,k} = (\\phi_i, J(\\phi_j, \\phi_k))`""" if not self.stored: - subs = self.subs + self.oceanic_basis.substitutions + if self.mk_subs: + subs = self.subs + self.oceanic_basis.substitutions + else: + subs = self.subs + args = ((i, j, k), self.ip.ip_jac, (self._phi(i), self._phi(j), self._phi(k)), subs) return self._integrate(subs, args) else: @@ -925,7 +1111,11 @@ def O(self, i, j, k): def C(self, i, j, k): """Function to compute the tensors holding the Jacobian inner products: :math:`C_{i,j,k} = (\\phi_i, J(\\phi_j,\\nabla^2 \\phi_k))`.""" if not self.stored: - subs = self.subs + self.oceanic_basis.substitutions + if self.mk_subs: + subs = self.subs + self.oceanic_basis.substitutions + else: + subs = self.subs + args = ((i, j, k), self.ip.ip_jac_lap, (self._phi(i), self._phi(j), self._phi(k)), subs) return self._integrate(subs, args) else: @@ -935,7 +1125,11 @@ def K(self, i, j): """Function to commpute the forcing of the ocean by the atmosphere: :math:`K_{i,j} = (\\phi_i, \\nabla^2 F_j)`.""" if self.connected_to_atmosphere: if not self.stored: - subs = self.subs + self.oceanic_basis.substitutions + self.atmospheric_basis.substitutions + if self.mk_subs: + subs = self.subs + self.oceanic_basis.substitutions + self.atmospheric_basis.substitutions + else: + subs = self.subs + args = ((i, j), self.iip.ip_lap, (self._phi(i), self._F(j)), subs) return self._integrate(subs, args) else: @@ -947,7 +1141,11 @@ def W(self, i, j): """Function to compute the short-wave radiative forcing of the ocean: :math:`W_{i,j} = (\\phi_i, F_j)`.""" if self.connected_to_atmosphere: if not self.stored: - subs = self.subs + self.oceanic_basis.substitutions + self.atmospheric_basis.substitutions + if self.mk_subs: + subs = self.subs + self.oceanic_basis.substitutions + self.atmospheric_basis.substitutions + else: + subs = self.subs + args = ((i, j), self.iip.symbolic_inner_product, (self._phi(i), self._F(j)), subs) return self._integrate(subs, args) else: @@ -958,7 +1156,11 @@ def W(self, i, j): def V(self, i, j, k, l, m): """Function to compute the :math:`T^4` temperature forcing from the ocean to the atmosphere: :math:`V_{i,j,k,l,m} = (\\phi_i, \\phi_j, \\phi_k, \\phi_l, \\phi_m)`.""" if not self.stored: - subs = self.subs + self.oceanic_basis.substitutions + if self.mk_subs: + subs = self.subs + self.oceanic_basis.substitutions + else: + subs = self.subs + args = ((i, j, k, l, m), self.ip.symbolic_inner_product, (self._phi(i), self._phi(j) * self._phi(k) * self._phi(l) * self._phi(m)), subs) return self._integrate(subs, args) else: @@ -967,7 +1169,11 @@ def V(self, i, j, k, l, m): def Z(self, i, j, k, l, m): """Function to compute the :math:`T^4` temperature forcing from the atmosphere to the ocean: :math:`Z_{i,j,k,l,m} = (\\phi_i, F_j, F_k, F_l, F_m)`.""" if not self.stored: - subs = self.subs + self.oceanic_basis.substitutions + self.atmospheric_basis.substitutions + if self.mk_subs: + subs = self.subs + self.oceanic_basis.substitutions + self.atmospheric_basis.substitutions + else: + subs = self.subs + args = ((i, j, k, l, m), self.ip.symbolic_inner_product, (self._phi(i), self._F(j) * self._F(k) * self._F(l) * self._F(m)), subs) return self._integrate(subs, args) else: @@ -1036,7 +1242,7 @@ class GroundSymbolicInnerProducts(GroundInnerProducts): """ def __init__(self, params=None, stored=True, inner_product_definition=None, interaction_inner_product_definition=None, - num_threads=None, quadrature=True, timeout=None, dynTinnerproducts=None, T4innerproducts=None): + num_threads=None, quadrature=True, timeout=None, dynTinnerproducts=None, T4innerproducts=None, return_symbolic=False, make_substitution=True): GroundInnerProducts.__init__(self) @@ -1086,7 +1292,16 @@ def __init__(self, params=None, stored=True, inner_product_definition=None, inte self.atmospheric_basis = None self.connected_to_atmosphere = False - self.subs = [(_n, self.n)] + self.return_symbolic = return_symbolic + if return_symbolic: + self.mk_subs = make_substitution + if self.mk_subs: + self.subs = [(self.n.symbol, self.n)] + else: + self.subs = None + else: + self.mk_subs = True + self.subs = [(self.n.symbol, self.n)] if inner_product_definition is None: self.ip = StandardSymbolicInnerProductDefinition() @@ -1139,35 +1354,47 @@ def connect_to_atmosphere(self, atmosphere_basis, num_threads=None, timeout=None num_threads = cpu_count() with Pool(max_workers=num_threads) as pool: - - subs = self.subs + self.ground_basis.substitutions + self.atmospheric_basis.substitutions + if self.mk_subs: + subs = self.subs + self.ground_basis.substitutions + self.atmospheric_basis.substitutions + else: + subs = self.subs natm = len(atmosphere_basis) - self._W = sp.zeros((self.ngr, natm), dtype=float, format='dok') - if self._T4 or self._dynamic_T: - self._Z = sp.zeros((self.ngr, natm, natm, natm, natm), dtype=float, format='dok') + if self.return_symbolic: + self._W = None + self._Z = None + else: + self._W = sp.zeros((self.ngr, natm), dtype=float, format='dok') + if self._T4 or self._dynamic_T: + self._Z = sp.zeros((self.ngr, natm, natm, natm, natm), dtype=float, format='dok') # W inner products args_list = [[(i, j), self.iip.symbolic_inner_product, (self._phi(i), self._F(j))] for i in range(self.ngr) for j in range(natm)] - _parallel_compute(pool, args_list, subs, self._W, timeout) + output = _parallel_compute(pool, args_list, subs, self._W, timeout, symbolic_int=not self.mk_subs) + if self.return_symbolic: + self._W = ImmutableSparseMatrix(self.ngr, natm, output) + else: + self._W = self._W.to_coo() if self._T4: # Z inner products args_list = [[(i, j, k, ell, m), self.ip.symbolic_inner_product, (self._phi(i), self._F(j) * self._F(k) * self._F(ell) * self._F(m))] for i in range(self.ngr) for j in range(natm) for k in range(j, natm) for ell in range(k, natm) for m in range(ell, natm)] - _parallel_compute(pool, args_list, subs, self._Z, timeout, permute=True) + output = _parallel_compute(pool, args_list, subs, self._Z, timeout, permute=True, symbolic_int=not self.mk_subs) elif self._dynamic_T: # Z inner products args_list = [[(i, 0, 0, 0, m), self.ip.symbolic_inner_product, (self._phi(i), self._F(0) * self._F(0) * self._F(0) * self._F(m))] for i in range(self.ngr) for m in range(natm)] - _parallel_compute(pool, args_list, subs, self._Z, timeout, permute=True) + output = _parallel_compute(pool, args_list, subs, self._Z, timeout, permute=True, symbolic_int=not self.mk_subs) - self._W = self._W.to_coo() if self._T4 or self._dynamic_T: - self._Z = self._Z.to_coo() + if self.return_symbolic: + self._Z = ImmutableSparseNDimArray(output, shape=(self.ngr, natm, natm, natm, natm)) + else: + self._Z = self._Z.to_coo() def compute_inner_products(self, num_threads=None, timeout=None): """Function computing and storing all the inner products at once. @@ -1183,38 +1410,51 @@ def compute_inner_products(self, num_threads=None, timeout=None): If `None` or `False`, no timeout occurs. Default to `None`. """ - - self._U = sp.zeros((self.ngr, self.ngr), dtype=float, format='dok') - if self._T4 or self._dynamic_T: - self._V = sp.zeros((self.ngr, self.ngr, self.ngr, self.ngr, self.ngr), dtype=float, format='dok') + if self.return_symbolic: + self._U = None + self._V = None + else: + self._U = sp.zeros((self.ngr, self.ngr), dtype=float, format='dok') + if self._T4 or self._dynamic_T: + self._V = sp.zeros((self.ngr, self.ngr, self.ngr, self.ngr, self.ngr), dtype=float, format='dok') if self.stored: if num_threads is None: num_threads = cpu_count() with Pool(max_workers=num_threads) as pool: - subs = self.subs + self.ground_basis.substitutions + if self.mk_subs: + subs = self.subs + self.ground_basis.substitutions + else: + subs = self.subs # U inner products args_list = [[(i, j), self.ip.symbolic_inner_product, (self._phi(i), self._phi(j))] for i in range(self.ngr) for j in range(self.ngr)] - _parallel_compute(pool, args_list, subs, self._U, timeout) + output = _parallel_compute(pool, args_list, subs, self._U, timeout, symbolic_int=not self.mk_subs) + if self.return_symbolic: + self._U = ImmutableSparseMatrix(self.ngr, self.ngr, output) + else: + self._U = self._U.to_coo() if self._T4: # V inner products args_list = [[(i, j, k, ell, m), self.ip.symbolic_inner_product, (self._phi(i), self._phi(j) * self._phi(k) * self._phi(ell) * self._phi(m))] for i in range(self.ngr) for j in range(self.ngr) for k in range(j, self.ngr) for ell in range(k, self.ngr) for m in range(ell, self.ngr)] - _parallel_compute(pool, args_list, subs, self._V, timeout, permute=True) + output = _parallel_compute(pool, args_list, subs, self._V, timeout, permute=True, symbolic_int=not self.mk_subs) + elif self._dynamic_T: # V inner products args_list = [[(i, 0, 0, 0, m), self.ip.symbolic_inner_product, (self._phi(i), self._phi(0) * self._phi(0) * self._phi(0) * self._phi(m))] for i in range(self.ngr) for m in range(self.ngr)] - _parallel_compute(pool, args_list, subs, self._V, timeout, permute=True) + output = _parallel_compute(pool, args_list, subs, self._V, timeout, permute=True, symbolic_int=not self.mk_subs) - self._U = self._U.to_coo() if self._T4 or self._dynamic_T: - self._V = self._V.to_coo() + if self.return_symbolic: + self._V = ImmutableSparseNDimArray(output, shape=(self.ngr, self.ngr, self.ngr, self.ngr, self.ngr)) + else: + self._V = self._V.to_coo() @property def ngr(self): @@ -1226,6 +1466,10 @@ def ngr(self): # !-----------------------------------------------------! def _integrate(self, subs, args): + if self.return_symbolic: + res = _apply(args) + return res[1] + if self.quadrature: res = _num_apply(args) return res[1] @@ -1254,7 +1498,11 @@ def M(self, i, j): def U(self, i, j): """Function to compute the inner products: :math:`U_{i,j} = (\\phi_i, \\phi_j)`.""" if not self.stored: - subs = self.subs + self.ground_basis.substitutions + if self.mk_subs: + subs = self.subs + self.ground_basis.substitutions + else: + subs = self.subs + args = ((i, j), self.ip.symbolic_inner_product, (self._phi(i), self._phi(j)), subs) return self._integrate(subs, args) else: @@ -1291,7 +1539,11 @@ def W(self, i, j): """Function to compute the short-wave radiative forcing of the ground: :math:`W_{i,j} = (\\phi_i, F_j)`.""" if self.connected_to_atmosphere: if not self.stored: - subs = self.subs + self.ground_basis.substitutions + self.atmospheric_basis.substitutions + if self.mk_subs: + subs = self.subs + self.ground_basis.substitutions + self.atmospheric_basis.substitutions + else: + subs = self.subs + args = ((i, j), self.iip.symbolic_inner_product, (self._phi(i), self._F(j)), subs) return self._integrate(subs, args) else: @@ -1302,7 +1554,11 @@ def W(self, i, j): def V(self, i, j, k, l, m): """Function to compute the :math:`T^4` temperature forcing from the ground to the atmosphere: :math:`V_{i,j,k,l,m} = (\\phi_i, \\phi_j, \\phi_k, \\phi_l, \\phi_m)`.""" if not self.stored: - subs = self.subs + self.ground_basis.substitutions + if self.mk_subs: + subs = self.subs + self.ground_basis.substitutions + else: + subs = self.subs + args = ((i, j, k, l, m), self.ip.symbolic_inner_product, (self._phi(i), self._phi(j) * self._phi(k) * self._phi(l) * self._phi(m)), subs) return self._integrate(subs, args) else: @@ -1311,7 +1567,11 @@ def V(self, i, j, k, l, m): def Z(self, i, j, k, l, m): """Function to compute the :math:`T^4` temperature forcing from the atmosphere to the ground: :math:`Z_{i,j,k,l,m} = (\\phi_i, F_j, F_k, F_l, F_m)`.""" if not self.stored: - subs = self.subs + self.ground_basis.substitutions + self.atmospheric_basis.substitutions + if self.mk_subs: + subs = self.subs + self.ground_basis.substitutions + self.atmospheric_basis.substitutions + else: + subs = self.subs + args = ((i, j, k, l, m), self.ip.symbolic_inner_product, (self._phi(i), self._F(j) * self._F(k) * self._F(l) * self._F(m)), subs) return self._integrate(subs, args) else: @@ -1325,9 +1585,13 @@ def _apply(ls): def _num_apply(ls): integrand = ls[1](*ls[2], integrand=True) - num_integrand = integrand[0].subs(ls[3]) + if ls[3] is not None: + num_integrand = integrand[0].subs(ls[3]) + else: + num_integrand = integrand[0] + func = lambdify((integrand[1][0], integrand[2][0]), num_integrand, 'numpy') - + try: a = integrand[2][1].subs(ls[3]) except: @@ -1369,9 +1633,14 @@ def _num_apply(ls): return ls[0], res[0] -def _parallel_compute(pool, args_list, subs, destination, timeout, permute=False): +def _parallel_compute(pool, args_list, subs, destination, timeout, permute=False, symbolic_int=False): + if destination is None: + return_dict = True + destination = dict() + else: + return_dict = False - if timeout is False: + if timeout is False or symbolic_int: timeout = None if timeout is not True: @@ -1382,7 +1651,18 @@ def _parallel_compute(pool, args_list, subs, destination, timeout, permute=False while True: try: res = next(results) - destination[res[0]] = float(res[1].subs(subs)) + if symbolic_int: + expr = res[1].simplify() + destination[res[0]] = expr + if permute: + i = res[0][0] + idx = res[0][1:] + perm_idx = multiset_permutations(idx) + for perm in perm_idx: + idx = [i] + perm + destination[tuple(idx)] = expr + else: + destination[res[0]] = float(res[1].subs(subs)) except StopIteration: break except TimeoutError: @@ -1413,6 +1693,9 @@ def _parallel_compute(pool, args_list, subs, destination, timeout, permute=False except StopIteration: break + if return_dict: + return destination + if __name__ == '__main__': from qgs.params.params import QgParams diff --git a/qgs/params/parameter.py b/qgs/params/parameter.py index 21f0f84..1917307 100644 --- a/qgs/params/parameter.py +++ b/qgs/params/parameter.py @@ -9,7 +9,8 @@ -------- >>> from qgs.params.params import ScaleParams - >>> from qgs.params.parameter import Parameter + >>> from qgs.params.parameter import Parameter, ParametersArray + >>> import numpy as np >>> # defining a scale object to help Parameter compute the nondimensionalization >>> sc = ScaleParams() >>> # creating a parameter initialized with a nondimensional value but returning a @@ -36,12 +37,438 @@ 2.1581898457499433e-06 >>> sigma.return_dimensional False + >>> # creating a parameters array initialized with a nondimensional values and returning + >>> # nondimensional ones when called + >>> s = ParametersArray(np.array([[0.1,0.2],[0.3,0.4]]), input_dimensional=False, scale_object=sc, units='[s^-1]', + ... description="atmosphere bottom friction coefficient") + >>> s + ArrayParameters([[0.1, 0.2], + [0.3, 0.4]], dtype=object) + >>> # dimensional values can also be retrieved with + >>> s.dimensional_values + array([[1.0320000000000001e-05, 2.0640000000000002e-05], + [3.096e-05, 4.1280000000000005e-05]], dtype=object) + >>> # you can also ask for the dimensional value of one particular value of the array + >>> s[0,0] + 0.1 + >>> s[0,0].dimensional_value + 1.0320000000000001e-05 Main class ---------- """ import warnings +import numpy as np +from fractions import Fraction + + +# TODO: Automatize warnings and errors +# TODO: Implement operations for arrays + +class ScalingParameter(float): + """Class of model's dimension parameter. + + Parameters + ---------- + value: float + Value of the parameter. + units: str, optional + The units of the provided value. Used to compute the conversion between dimensional and nondimensional + value. Should be specified by joining atoms like `'[unit^power]'`, e.g '`[m^2][s^-2][Pa^-2]'`. + Empty by default. + description: str, optional + String describing the parameter. + dimensional: bool, optional + Indicate if the value of the parameter is dimensional or not. Default to `True`. + symbol: ~sympy.core.symbol.Symbol, optional + A `Sympy`_ symbol to represent the parameter in symbolic expressions. + symbolic_expression: ~sympy.core.expr.Expr, optional + A `Sympy`_ expression to represent a relationship to other parameters. + + Notes + ----- + Parameter is immutable. Once instantiated, it cannot be altered. To create a new parameter, one must + re-instantiate it. + + .. _Sympy: https://www.sympy.org/ + """ + + def __new__(cls, value, units="", description="", dimensional=False, symbol=None, symbolic_expression=None): + + f = float.__new__(cls, value) + f._dimensional = dimensional + f._units = units + f._description = description + f._symbol = symbol + f._symbolic_expression = symbolic_expression + + return f + + @property + def symbol(self): + """~sympy.core.symbol.Symbol: Returns the symbol of the parameter.""" + return self._symbol + + @property + def symbolic_expression(self): + """~sympy.core.expr.Expr: Returns the symbolic expression of the parameter.""" + if self._symbolic_expression is None and self._symbol is not None: + return self._symbol + else: + return self._symbolic_expression + + @property + def dimensional(self): + """bool: Indicate if the returned value is dimensional or not.""" + return self._dimensional + + @property + def units(self): + """str: The units of the dimensional value.""" + return self._units + + @property + def description(self): + """str: Description of the parameter.""" + return self._description + + def __add__(self, other): + + res = float(self) + other + if isinstance(other, (Parameter, ScalingParameter)): + if self.units != other.units: + raise ArithmeticError("ScalingParameter class: Impossible to add two parameters with different units.") + if self.symbolic_expression is None: + if other.symbolic_expression is None: + if self.symbol is not None and other.symbol is not None: + expr = self.symbol + other.symbol + else: + expr = None + descr = self.description + " + " + other.description + else: + if self.symbol is not None: + expr = self.symbol + (other.symbolic_expression) + descr = self.description + " + (" + other.description + ")" + else: + expr = None + descr = self.description + " + " + other.description + else: + if other.symbolic_expression is None: + if other.symbol is not None: + expr = (self.symbolic_expression) + other.symbol + descr = "(" + self.description + ") + " + other.description + else: + expr = None + descr = self.description + " + " + other.description + else: + expr = (self.symbolic_expression) + (other.symbolic_expression) + descr = "(" + self.description + ") + (" + other.description + ")" + + if isinstance(other, Parameter): + return Parameter(res, input_dimensional=other.input_dimensional, + return_dimensional=other.return_dimensional, scale_object=other._scale_object, + description=descr, units=self.units, symbol=None, symbolic_expression=expr) + + else: + return ScalingParameter(res, description=descr, units=self.units, symbol=None, symbolic_expression=expr) + else: + try: + if self.symbol is not None: + expr = self.symbol + other + descr = self.description + " + " + str(other) + elif self.symbolic_expression is not None: + expr = (self.symbolic_expression) + other + descr = "(" + self.description + ") + " + str(other) + else: + expr = None + descr = self.description + " + " + str(other) + return ScalingParameter(res, description=descr, units=self.units, symbol=None, symbolic_expression=expr) + except: + return res + + def __radd__(self, other): + return self.__add__(other) + + def __sub__(self, other): + + res = float(self) - other + if isinstance(other, (Parameter, ScalingParameter)): + if self.units != other.units: + raise ArithmeticError("ScalingParameter class: Impossible to subtract two parameters with different units.") + if self.symbolic_expression is None: + if other.symbolic_expression is None: + if self.symbol is not None and other.symbol is not None: + expr = self.symbol - other.symbol + else: + expr = None + descr = self.description + " - " + other.description + else: + if self.symbol is not None: + expr = self.symbol - (other.symbolic_expression) + descr = self.description + " - (" + other.description + ")" + else: + expr = None + descr = self.description + " - " + other.description + else: + if other.symbolic_expression is None: + if other.symbol is not None: + expr = (self.symbolic_expression) - other.symbol + descr = "(" + self.description + ") - " + other.description + else: + expr = None + descr = self.description + " - " + other.description + else: + expr = (self.symbolic_expression) - (other.symbolic_expression) + descr = "(" + self.description + ") - (" + other.description + ")" + + if isinstance(other, Parameter): + return Parameter(res, input_dimensional=other.input_dimensional, + return_dimensional=other.return_dimensional, scale_object=other._scale_object, + description=descr, units=self.units, symbol=None, symbolic_expression=expr) + + else: + return ScalingParameter(res, description=descr, units=self.units, symbol=None, symbolic_expression=expr) + else: + try: + if self.symbol is not None: + expr = self.symbol - other + descr = self.description + " - " + str(other) + elif self.symbolic_expression is not None: + expr = (self.symbolic_expression) - other + descr = "(" + self.description + ") - " + str(other) + else: + expr = None + descr = self.description + " - " + str(other) + return ScalingParameter(res, description=descr, units=self.units, symbol=None, symbolic_expression=expr) + except: + return res + + def __rsub__(self, other): + try: + res = other - float(self) + if self.symbol is not None: + expr = other - self.symbol + descr = str(other) + " - " + self.description + elif self.symbolic_expression is not None: + expr = other - (self.symbolic_expression) + descr = str(other) + " - (" + self.description + ")" + else: + expr = None + descr = str(other) + " - " + self.description + return ScalingParameter(res, description=descr, units=self.units, symbol=None, symbolic_expression=expr) + except: + return res + + def __mul__(self, other): + + res = float(self) * other + if isinstance(other, (Parameter, ScalingParameter)): + units = _combine_units(self.units, other.units, '+') + + if self.symbolic_expression is None: + if other.symbolic_expression is None: + if self.symbol is not None and other.symbol is not None: + expr = self.symbol * other.symbol + else: + expr = None + descr = self.description + " * " + other.description + else: + if self.symbol is not None: + expr = self.symbol * (other.symbolic_expression) + descr = self.description + " * (" + other.description + ")" + else: + expr = None + descr = self.description + " * " + other.description + else: + if other.symbolic_expression is None: + if other.symbol is not None: + expr = (self.symbolic_expression) * other.symbol + descr = "(" + self.description + ") * " + other.description + else: + expr = None + descr = self.description + " * " + other.description + else: + expr = (self.symbolic_expression) * (other.symbolic_expression) + descr = "(" + self.description + ") * (" + other.description + ")" + + if isinstance(other, Parameter): + return Parameter(res, input_dimensional=other.input_dimensional, return_dimensional=other.return_dimensional, + scale_object=other._scale_object, description=descr, + units=units, symbol=None, symbolic_expression=expr) + + else: + return ScalingParameter(res, description=descr, units=units, symbol=None, symbolic_expression=expr) + else: + try: + if self.symbol is not None: + expr = self.symbol * other + descr = self.description + " * " + str(other) + elif self.symbolic_expression is not None: + expr = (self.symbolic_expression) * other + descr = "(" + self.description + ") * " + str(other) + else: + expr = None + descr = self.description + " * " + str(other) + return ScalingParameter(res, description=descr, units=self.units, symbol=None, symbolic_expression=expr) + except: + return res + + def __rmul__(self, other): + return self.__mul__(other) + + def __truediv__(self, other): + + res = float(self) / other + if isinstance(other, (ScalingParameter, Parameter)): + units = _combine_units(self.units, other.units, '-') + if self.symbolic_expression is None: + if other.symbolic_expression is None: + if self.symbol is not None and other.symbol is not None: + expr = self.symbol / other.symbol + else: + expr = None + descr = self.description + " / " + other.description + else: + if self.symbol is not None: + expr = self.symbol / (other.symbolic_expression) + descr = self.description + " / (" + other.description + ")" + else: + expr = None + descr = self.description + " / " + other.description + else: + if other.symbolic_expression is None: + if other.symbol is not None: + expr = (self.symbolic_expression) / other.symbol + descr = "(" + self.description + ") / " + other.description + else: + expr = None + descr = self.description + " / " + other.description + else: + expr = (self.symbolic_expression) / (other.symbolic_expression) + descr = "(" + self.description + ") / (" + other.description + ")" + + if isinstance(other, Parameter): + return Parameter(res, input_dimensional=other.input_dimensional, return_dimensional=other.return_dimensional, + scale_object=other._scale_object, description=descr, + units=units, symbol=None, symbolic_expression=expr) + else: + return ScalingParameter(res, description=descr, units=units, symbol=None, symbolic_expression=expr) + else: + try: + if self.symbol is not None: + expr = self.symbol / other + descr = self.description + " / " + str(other) + elif self.symbolic_expression is not None: + expr = (self.symbolic_expression) / other + descr = "(" + self.description + ") / " + str(other) + else: + expr = None + descr = self.description + " / " + str(other) + return ScalingParameter(res, description=descr, units=self.units, symbol=None, symbolic_expression=expr) + except: + return res + + def __rtruediv__(self, other): + res = other / float(self) + try: + if self.symbol is not None: + expr = other / self.symbol + descr = str(other) + " / " + self.description + elif self.symbolic_expression is not None: + expr = other / (self.symbolic_expression) + descr = str(other) + " / (" + self.description + ")" + else: + expr = None + descr = str(other) + " / " + self.description + return ScalingParameter(res, description=descr, units=self.units, symbol=None, symbolic_expression=expr) + except: + return res + + def __pow__(self, power, modulo=None): + + if modulo is not None: + raise NotImplemented('ScalingParameter class: Modular exponentiation not implemented') + + res = float(self) ** power + if int(power) == power: + + ul = self.units.split('][') + ul[0] = ul[0][1:] + ul[-1] = ul[-1][:-1] + + usl = list() + for us in ul: + up = us.split('^') + if len(up) == 1: + up.append("1") + + usl.append(tuple(up)) + + units_elements = list() + for us in usl: + units_elements.append(list((us[0], str(int(us[1]) * power)))) + + units = list() + for us in units_elements: + if us is not None: + if int(us[1]) != 1: + units.append("[" + us[0] + "^" + us[1] + "]") + else: + units.append("[" + us[0] + "]") + units = "".join(units) + + if self.symbolic_expression is not None: + expr = (self.symbolic_expression) ** power + descr = "(" + self.description + ") to the power "+str(power) + elif self.symbol is not None: + expr = self.symbol ** power + descr = self.description + " to the power "+str(power) + else: + expr = None + descr = self.description + " to the power "+str(power) + + else: + power_fraction = Fraction(power) + ul = self.units.split('][') + ul[0] = ul[0][1:] + ul[-1] = ul[-1][:-1] + + usl = list() + for us in ul: + up = us.split('^') + if len(up) == 1: + up.append("1") + + usl.append(tuple(up)) + + units_elements = list() + for us in usl: + new_power = int(us[1]) * power_fraction.numerator / power_fraction.denominator + if int(new_power) == new_power: + units_elements.append(list((us[0], str(new_power)))) + else: + raise ArithmeticError("ScalingParameter class: Only support integer exponent in units") + + units = list() + for us in units_elements: + if us is not None: + if int(us[1]) != 1: + units.append("[" + us[0] + "^" + us[1] + "]") + else: + units.append("[" + us[0] + "]") + units = "".join(units) + if self.symbolic_expression is not None: + expr = (self.symbolic_expression) ** power + descr = "(" + self.description + ") to the power " + str(power) + elif self.symbol is not None: + expr = self.symbol ** power + descr = self.description + " to the power " + str(power) + else: + expr = None + descr = self.description + " to the power " + str(power) + + return ScalingParameter(res, description=descr, units=units, symbol=None, symbolic_expression=expr) class Parameter(float): @@ -62,6 +489,10 @@ class Parameter(float): `None` by default. If `None`, cannot transform between dimensional and nondimentional value. description: str, optional String describing the parameter. + symbol: ~sympy.core.symbol.Symbol, optional + A `Sympy`_ symbol to represent the parameter in symbolic expressions. + symbolic_expression: ~sympy.core.expr.Expr, optional + A `Sympy`_ expression to represent a relationship to other parameters. return_dimensional: bool, optional Defined if the value returned by the parameter is dimensional or not. Default to `False`. @@ -74,10 +505,11 @@ class Parameter(float): -------- If no scale_object argument is provided, cannot transform between the dimensional and nondimentional value ! + .. _Sympy: https://www.sympy.org/ """ def __new__(cls, value, input_dimensional=True, units="", scale_object=None, description="", - return_dimensional=False): + symbol=None, return_dimensional=False, symbolic_expression=None): no_scale = False @@ -98,7 +530,14 @@ def __new__(cls, value, input_dimensional=True, units="", scale_object=None, des evalue = value no_scale = True else: - evalue = value * cls._conversion_factor(units, scale_object) + try: + evalue = value * cls._conversion_factor(units, scale_object) + except: + print(description) + print(symbol) + print(units) + print(cls._conversion_factor(units, scale_object)) + print(scale_object) else: evalue = value @@ -112,6 +551,8 @@ def __new__(cls, value, input_dimensional=True, units="", scale_object=None, des f._units = units f._scale_object = scale_object f._description = description + f._symbol = symbol + f._symbolic_expression = symbolic_expression return f @@ -131,6 +572,19 @@ def nondimensional_value(self): else: return self + @property + def symbol(self): + """~sympy.core.symbol.Symbol: Returns the symbol of the parameter.""" + return self._symbol + + @property + def symbolic_expression(self): + """~sympy.core.expr.Expr: Returns the symbolic expression of the parameter.""" + if self._symbolic_expression is None and self._symbol is not None: + return self._symbol + else: + return self._symbolic_expression + @property def input_dimensional(self): """bool: Indicate if the provided value is dimensional or not.""" @@ -180,4 +634,715 @@ def _nondimensionalization(self): else: return self._conversion_factor(self._units, self._scale_object) + def __add__(self, other): + + res = float(self) + other + if isinstance(other, (Parameter, ScalingParameter)): + if isinstance(other, Parameter) and self.return_dimensional != other.return_dimensional: + raise ArithmeticError("Parameter class: Impossible to subtract a dimensional parameter with a non-dimensional one.") + if self.units != other.units: + raise ArithmeticError("Parameter class: Impossible to add two parameters with different units.") + + if self.symbolic_expression is None: + if other.symbolic_expression is None: + if self.symbol is not None and other.symbol is not None: + expr = self.symbol + other.symbol + else: + expr = None + descr = self.description + " + " + other.description + else: + if self.symbol is not None: + expr = self.symbol + (other.symbolic_expression) + descr = self.description + " + (" + other.description + ")" + else: + expr = None + descr = self.description + " + " + other.description + else: + if other.symbolic_expression is None: + if other.symbol is not None: + expr = (self.symbolic_expression) + other.symbol + descr = "(" + self.description + ") + " + other.description + else: + expr = None + descr = self.description + " + " + other.description + else: + expr = (self.symbolic_expression) + (other.symbolic_expression) + descr = "(" + self.description + ") + (" + other.description + ")" + + return Parameter(res, input_dimensional=self.return_dimensional, + return_dimensional=self.return_dimensional, scale_object=self._scale_object, + description=descr, units=self.units, symbol=None, symbolic_expression=expr) + else: + try: + if self.symbol is not None: + expr = self.symbol + other + descr = self.description + " + " + str(other) + elif self.symbolic_expression is not None: + expr = (self.symbolic_expression) + other + descr = "(" + self.description + ") + " + str(other) + else: + expr = None + descr = self.description + " + " + str(other) + return Parameter(res, input_dimensional=self.return_dimensional, + return_dimensional=self.return_dimensional, scale_object=self._scale_object, + description=descr, units=self.units, symbol=None, symbolic_expression=expr) + except: + return res + + def __radd__(self, other): + return self.__add__(other) + def __sub__(self, other): + + res = float(self) - other + if isinstance(other, (Parameter, ScalingParameter)): + if isinstance(other, Parameter) and self.return_dimensional != other.return_dimensional: + raise ArithmeticError("Parameter class: Impossible to subtract a dimensional parameter with a non-dimensional one.") + if self.units != other.units: + raise ArithmeticError("Parameter class: Impossible to subtract two parameters with different units.") + if self.symbolic_expression is None: + if other.symbolic_expression is None: + if self.symbol is not None and other.symbol is not None: + expr = self.symbol - other.symbol + else: + expr = None + descr = self.description + " - " + other.description + else: + if self.symbol is not None: + expr = self.symbol - (other.symbolic_expression) + descr = self.description + " - (" + other.description + ")" + else: + expr = None + descr = self.description + " - " + other.description + else: + if other.symbolic_expression is None: + if other.symbol is not None: + expr = (self.symbolic_expression) - other.symbol + descr = "(" + self.description + ") - " + other.description + else: + expr = None + descr = self.description + " - " + other.description + else: + expr = (self.symbolic_expression) - (other.symbolic_expression) + descr = "(" + self.description + ") - (" + other.description + ")" + + return Parameter(res, input_dimensional=self.return_dimensional, + return_dimensional=self.return_dimensional, scale_object=self._scale_object, + description=descr, units=self.units, symbol=None, symbolic_expression=expr) + else: + try: + if self.symbol is not None: + expr = self.symbol - other + descr = self.description + " - " + str(other) + elif self.symbolic_expression is not None: + expr = (self.symbolic_expression) - other + descr = "(" + self.description + ") - " + str(other) + else: + expr = None + descr = self.description + " - " + str(other) + return Parameter(res, input_dimensional=self.return_dimensional, + return_dimensional=self.return_dimensional, scale_object=self._scale_object, + description=descr, units=self.units, symbol=None, symbolic_expression=expr) + except: + return res + + def __rsub__(self, other): + res = other - float(self) + try: + if self.symbol is not None: + expr = other - self.symbol + descr = str(other) + " - " + self.description + elif self.symbolic_expression is not None: + expr = other - (self.symbolic_expression) + descr = str(other) + " - (" + self.description + ")" + else: + expr = None + descr = str(other) + " - " + self.description + return Parameter(res, input_dimensional=self.return_dimensional, + return_dimensional=self.return_dimensional, scale_object=self._scale_object, + description=descr, units=self.units, symbol=None, symbolic_expression=expr) + except: + return res + + def __mul__(self, other): + + res = float(self) * other + if isinstance(other, (Parameter, ScalingParameter)): + if hasattr(other, "units"): + units = _combine_units(self.units, other.units, '+') + else: + units = "" + + if self.symbolic_expression is None: + if other.symbolic_expression is None: + if self.symbol is not None and other.symbol is not None: + expr = self.symbol * other.symbol + else: + expr = None + descr = self.description + " * " + other.description + else: + if self.symbol is not None: + expr = self.symbol * (other.symbolic_expression) + descr = self.description + " * (" + other.description + ")" + else: + expr = None + descr = self.description + " * " + other.description + else: + if other.symbolic_expression is None: + if other.symbol is not None: + expr = (self.symbolic_expression) * other.symbol + descr = "(" + self.description + ") * " + other.description + else: + expr = None + descr = self.description + " * " + other.description + else: + expr = (self.symbolic_expression) * (other.symbolic_expression) + descr = "(" + self.description + ") * (" + other.description + ")" + + return Parameter(res, input_dimensional=self.return_dimensional, return_dimensional=self.return_dimensional, + scale_object=self._scale_object, description=descr, units=units, symbol=None, + symbolic_expression=expr) + else: + try: + if self.symbol is not None: + expr = self.symbol * other + descr = self.description + " * " + str(other) + elif self.symbolic_expression is not None: + expr = (self.symbolic_expression) * other + descr = "(" + self.description + ") * " + str(other) + else: + expr = None + descr = self.description + " * " + str(other) + return Parameter(res, input_dimensional=self.return_dimensional, return_dimensional=self.return_dimensional, + scale_object=self._scale_object, description=descr, units=self.units, symbol=None, + symbolic_expression=expr) + except: + return res + + def __rmul__(self, other): + return self.__mul__(other) + + def __truediv__(self, other): + + res = float(self) / other + if isinstance(other, (ScalingParameter, Parameter)): + units = _combine_units(self.units, other.units, '-') + if self.symbolic_expression is None: + if other.symbolic_expression is None: + if self.symbol is not None and other.symbol is not None: + expr = self.symbol / other.symbol + else: + expr = None + descr = self.description + " / " + other.description + else: + if self.symbol is not None: + expr = self.symbol / (other.symbolic_expression) + descr = self.description + " / (" + other.description + ")" + else: + expr = None + descr = self.description + " / " + other.description + else: + if other.symbolic_expression is None: + if other.symbol is not None: + expr = (self.symbolic_expression) / other.symbol + descr = "(" + self.description + ") / " + other.description + else: + expr = None + descr = self.description + " / " + other.description + else: + expr = (self.symbolic_expression) / (other.symbolic_expression) + descr = "(" + self.description + ") / (" + other.description + ")" + + return Parameter(res, input_dimensional=self.return_dimensional, return_dimensional=self.return_dimensional, + scale_object=self._scale_object, description=descr, units=units, symbol=None, + symbolic_expression=expr) + else: + try: + if self.symbol is not None: + expr = self.symbol / other + descr = self.description + " / " + str(other) + elif self.symbolic_expression is not None: + expr = (self.symbolic_expression) / other + descr = "(" + self.description + ") / " + str(other) + else: + expr = None + descr = self.description + " / " + str(other) + return Parameter(res, input_dimensional=self.return_dimensional, return_dimensional=self.return_dimensional, + scale_object=self._scale_object, description=descr, units=self.units, symbol=None, + symbolic_expression=expr) + except: + return res + + def __rtruediv__(self, other): + res = other / float(self) + try: + if self.symbol is not None: + expr = other / self.symbol + descr = str(other) + " / " + self.description + elif self.symbolic_expression is not None: + expr = other / (self.symbolic_expression) + descr = str(other) + " / (" + self.description + ")" + else: + expr = None + descr = str(other) + " / " + self.description + return Parameter(res, input_dimensional=self.return_dimensional, + return_dimensional=self.return_dimensional, scale_object=self._scale_object, + description=descr, units=self.units, symbol=None, symbolic_expression=expr) + except: + return res + + def __pow__(self, power, modulo=None): + + if modulo is not None: + raise NotImplemented('Parameter class: Modular exponentiation not implemented') + + res = float(self) ** power + if int(power) == power: + + ul = self.units.split('][') + ul[0] = ul[0][1:] + ul[-1] = ul[-1][:-1] + + usl = list() + for us in ul: + up = us.split('^') + if len(up) == 1: + up.append("1") + + usl.append(tuple(up)) + + units_elements = list() + for us in usl: + units_elements.append(list((us[0], str(int(us[1]) * power)))) + + units = list() + for us in units_elements: + if us is not None: + if int(us[1]) != 1: + units.append("[" + us[0] + "^" + us[1] + "]") + else: + units.append("[" + us[0] + "]") + units = "".join(units) + + if self.symbolic_expression is not None: + expr = (self.symbolic_expression) ** power + descr = "(" + self.description + ") to the power "+str(power) + elif self.symbol is not None: + expr = self.symbol ** power + descr = self.description + " to the power "+str(power) + else: + expr = None + descr = self.description + " to the power "+str(power) + + else: + power_fraction = Fraction(power) + ul = self.units.split('][') + ul[0] = ul[0][1:] + ul[-1] = ul[-1][:-1] + + usl = list() + for us in ul: + up = us.split('^') + if len(up) == 1: + up.append("1") + + usl.append(tuple(up)) + + units_elements = list() + for us in usl: + new_power = int(us[1]) * power_fraction.numerator / power_fraction.denominator + if int(new_power) == new_power: + units_elements.append(list((us[0], str(int(new_power))))) + else: + raise ArithmeticError("Parameter class: Only support integer exponent in units") + + units = list() + for us in units_elements: + if us is not None: + if int(us[1]) != 1: + units.append("[" + us[0] + "^" + us[1] + "]") + else: + units.append("[" + us[0] + "]") + units = "".join(units) + if self.symbolic_expression is not None: + expr = (self.symbolic_expression) ** power + descr = "(" + self.description + ") to the power "+str(power) + elif self.symbol is not None: + expr = self.symbol ** power + descr = self.description + " to the power "+str(power) + else: + expr = None + descr = self.description + " to the power "+str(power) + + return Parameter(res, input_dimensional=self.return_dimensional, return_dimensional=self.return_dimensional, + description=descr, units=units, scale_object=self._scale_object, symbol=None, + symbolic_expression=expr) + + +class ParametersArray(np.ndarray): + """Base class of model's array of parameters. + + Parameters + ---------- + values: list(float) or ~numpy.ndarray(float) or list(Parameter) or ~numpy.ndarray(Parameter) or list(ScalingParameter) or ~numpy.ndarray(ScalingParameter) + Values of the parameter array. + input_dimensional: bool, optional + Specify whether the value provided is dimensional or not. Default to `True`. + units: str, optional + The units of the provided value. Used to compute the conversion between dimensional and nondimensional + value. Should be specified by joining atoms like `'[unit^power]'`, e.g '`[m^2][s^-2][Pa^-2]'`. + Empty by default. + scale_object: ScaleParams, optional + A scale parameters object to compute the conversion between dimensional and nondimensional value. + `None` by default. If `None`, cannot transform between dimensional and nondimentional value. + description: str or list(str) or array(str), optional + String or an iterable of strings, describing the parameters. + If an iterable, should have the same length or shape as `values`. + symbols ~sympy.core.symbol.Symbol or list(~sympy.core.symbol.Symbol) or ~numpy.ndarray(~sympy.core.symbol.Symbol), optional + A `Sympy`_ symbol or an iterable of symbols, to represent the parameters in symbolic expressions. + If an iterable, should have the same length or shape as `values`. + symbolic_expressions: ~sympy.core.expr.Expr or list(~sympy.core.expr.Expr) or ~numpy.ndarray(~sympy.core.expr.Expr), optional + A `Sympy`_ expression or an iterable of expressions, to represent a relationship to other parameters. + If an iterable, should have the same length or shape as `values`. + return_dimensional: bool, optional + Defined if the value returned by the parameter is dimensional or not. Default to `False`. + + Warnings + -------- + If no scale_object argument is provided, cannot transform between the dimensional and nondimensional value ! + + .. _Sympy: https://www.sympy.org/ + """ + + def __new__(cls, values, input_dimensional=True, units="", scale_object=None, description="", + symbols=None, symbolic_expressions=None, return_dimensional=False): + + if isinstance(values, (tuple, list)): + new_arr = np.empty(len(values), dtype=object) + for i, val in enumerate(values): + if isinstance(description, (tuple, list, np.ndarray)): + descr = description[i] + else: + descr = description + if isinstance(symbols, (tuple, list, np.ndarray)): + sy = symbols[i] + else: + sy = symbols + if isinstance(symbolic_expressions, (tuple, list, np.ndarray)): + expr = symbolic_expressions[i] + else: + expr = symbolic_expressions + new_arr[i] = Parameter(val, input_dimensional=input_dimensional, units=units, scale_object=scale_object, description=descr, + return_dimensional=return_dimensional, symbol=sy, symbolic_expression=expr) + else: + if isinstance(values.flatten()[0], (Parameter, ScalingParameter)): + new_arr = values.copy() + else: + new_arr = np.empty_like(values, dtype=object) + for idx in np.ndindex(values.shape): + if isinstance(description, np.ndarray): + descr = description[idx] + else: + descr = description + if isinstance(symbols, np.ndarray): + sy = symbols[idx] + else: + sy = symbols + if isinstance(symbolic_expressions, np.ndarray): + expr = symbolic_expressions[idx] + else: + expr = symbolic_expressions + new_arr[idx] = Parameter(values[idx], input_dimensional=input_dimensional, units=units, scale_object=scale_object, description=descr, + return_dimensional=return_dimensional, symbol=sy, symbolic_expression=expr) + arr = np.asarray(new_arr).view(cls) + arr._input_dimensional = input_dimensional + arr._return_dimensional = return_dimensional + arr._units = units + arr._scale_object = scale_object + + return arr + + def __array_finalize__(self, arr): + + if arr is None: + return + + self._input_dimensional = getattr(arr, '_input_dimensional', True) + self._units = getattr(arr, '_units', "") + self._return_dimensional = getattr(arr, '_return_dimensional', False) + self._scale_object = getattr(arr, '_scale_object', None) + + @property + def dimensional_values(self): + """float: Returns the dimensional value.""" + if self._return_dimensional: + return self + else: + return np.array(self / self._nondimensionalization) + + @property + def nondimensional_values(self): + """float: Returns the nondimensional value.""" + if self._return_dimensional: + return np.array(self * self._nondimensionalization) + else: + return self + + @property + def symbols(self): + """~numpy.ndarray(~sympy.core.symbol.Symbol): Returns the symbol of the parameters in the array.""" + symbols = np.empty(self.shape, dtype=object) + for idx in np.ndindex(self.shape): + symbols[idx] = self[idx].symbol + return symbols + + @property + def symbolic_expressions(self): + """~numpy.ndarray(~sympy.core.expr.Expr): Returns the symbolic expressions of the parameters in the array.""" + symbolic_expressions = np.empty(self.shape, dtype=object) + for idx in np.ndindex(self.shape): + symbolic_expressions[idx] = self[idx].symbolic_expression + return symbolic_expressions + + @property + def input_dimensional(self): + """bool: Indicate if the provided value is dimensional or not.""" + return self._input_dimensional + + @property + def return_dimensional(self): + """bool: Indicate if the returned value is dimensional or not.""" + return self._return_dimensional + + @classmethod + def _conversion_factor(cls, units, scale_object): + factor = 1. + + ul = units.split('][') + ul[0] = ul[0][1:] + ul[-1] = ul[-1][:-1] + + for us in ul: + up = us.split('^') + if len(up) == 1: + up.append("1") + + if up[0] == 'm': + factor *= scale_object.L ** (-int(up[1])) + elif up[0] == 's': + factor *= scale_object.f0 ** (int(up[1])) + elif up[0] == 'Pa': + factor *= scale_object.deltap ** (-int(up[1])) + + return factor + + @property + def units(self): + """str: The units of the dimensional value.""" + return self._units + + @property + def descriptions(self): + """~numpy.ndarray(str): Description of the parameters in the array.""" + descr = np.empty(self.shape, dtype=object) + for idx in np.ndindex(self.shape): + descr[idx] = self[idx].description + return descr + + @property + def _nondimensionalization(self): + if self._scale_object is None: + return 1. + else: + return self._conversion_factor(self._units, self._scale_object) + + def __add__(self, other): + if isinstance(other, (Parameter, ScalingParameter, float, int)): + res = np.empty(self.shape, dtype=object) + for idx in np.ndindex(self.shape): + res[idx] = self[idx] + other + item = res[idx] + return ParametersArray(res, input_dimensional=item.return_dimensional, return_dimensional=item.return_dimensional, + units=item.units, scale_object=self._scale_object) + elif isinstance(other, ParametersArray): + if other.shape == self.shape: # Does not do broadcast + res = np.empty(self.shape, dtype=object) + for idx in np.ndindex(self.shape): + res[idx] = self[idx] + other[idx] + item = res[idx] + return ParametersArray(res, input_dimensional=item.return_dimensional, return_dimensional=item.return_dimensional, + units=item.units, scale_object=self._scale_object) + else: + return self + other + else: + return self + other + + def __radd__(self, other): + return self.__add__(other) + + def __sub__(self, other): + if isinstance(other, (Parameter, ScalingParameter, float, int)): + res = np.empty(self.shape, dtype=object) + for idx in np.ndindex(self.shape): + res[idx] = self[idx] - other + item = res[idx] + return ParametersArray(res, input_dimensional=item.return_dimensional, return_dimensional=item.return_dimensional, + units=item.units, scale_object=self._scale_object) + elif isinstance(other, ParametersArray): + if other.shape == self.shape: # Does not do broadcast + res = np.empty(self.shape, dtype=object) + for idx in np.ndindex(self.shape): + res[idx] = self[idx] - other[idx] + item = res[idx] + return ParametersArray(res, input_dimensional=item.return_dimensional, return_dimensional=item.return_dimensional, + units=item.units, scale_object=self._scale_object) + else: + return self - other + else: + return self - other + + def __rsub__(self, other): + if isinstance(other, (Parameter, ScalingParameter, float, int)): + res = np.empty(self.shape, dtype=object) + for idx in np.ndindex(self.shape): + res[idx] = other - self[idx] + item = res[idx] + return ParametersArray(res, input_dimensional=item.return_dimensional, return_dimensional=item.return_dimensional, + units=item.units, scale_object=self._scale_object) + elif isinstance(other, ParametersArray): + if other.shape == self.shape: # Does not do broadcast + res = np.empty(self.shape, dtype=object) + for idx in np.ndindex(self.shape): + res[idx] = other - self[idx] + item = res[idx] + return ParametersArray(res, input_dimensional=item.return_dimensional, return_dimensional=item.return_dimensional, + units=item.units, scale_object=self._scale_object) + else: + return other - self + else: + return other - self + + def __mul__(self, other): + if isinstance(other, (Parameter, ScalingParameter, float, int)): + res = np.empty(self.shape, dtype=object) + for idx in np.ndindex(self.shape): + res[idx] = self[idx] * other + item = res[idx] + return ParametersArray(res, input_dimensional=item.return_dimensional, return_dimensional=item.return_dimensional, + units=item.units, scale_object=self._scale_object) + elif isinstance(other, ParametersArray): + if other.shape == self.shape: # Does not do broadcast + res = np.empty(self.shape, dtype=object) + for idx in np.ndindex(self.shape): + res[idx] = self[idx] * other[idx] + item = res[idx] + return ParametersArray(res, input_dimensional=item.return_dimensional, return_dimensional=item.return_dimensional, + units=item.units, scale_object=self._scale_object) + else: + return self * other + else: + return self * other + + def __rmul__(self, other): + return self.__mul__(other) + + def __truediv__(self, other): + if isinstance(other, (Parameter, ScalingParameter, float, int)): + res = np.empty(self.shape, dtype=object) + for idx in np.ndindex(self.shape): + res[idx] = self[idx] / other + item = res[idx] + return ParametersArray(res, input_dimensional=item.return_dimensional, return_dimensional=item.return_dimensional, + units=item.units, scale_object=self._scale_object) + elif isinstance(other, ParametersArray): + if other.shape == self.shape: # Does not do broadcast + res = np.empty(self.shape, dtype=object) + for idx in np.ndindex(self.shape): + res[idx] = self[idx] / other[idx] + item = res[idx] + return ParametersArray(res, input_dimensional=item.return_dimensional, return_dimensional=item.return_dimensional, + units=item.units, scale_object=self._scale_object) + else: + return self / other + else: + return self / other + + def __rtruediv__(self, other): + if isinstance(other, (Parameter, ScalingParameter, float, int)): + res = np.empty(self.shape, dtype=object) + for idx in np.ndindex(self.shape): + res[idx] = other / self[idx] + item = res[idx] + return ParametersArray(res, input_dimensional=item.return_dimensional, return_dimensional=item.return_dimensional, + units=item.units, scale_object=self._scale_object) + elif isinstance(other, ParametersArray): + if other.shape == self.shape: # Does not do broadcast + res = np.empty(self.shape, dtype=object) + for idx in np.ndindex(self.shape): + res[idx] = other / self[idx] + item = res[idx] + return ParametersArray(res, input_dimensional=item.return_dimensional, return_dimensional=item.return_dimensional, + units=item.units, scale_object=self._scale_object) + else: + return other / self + else: + return other / self + + +def _combine_units(units1, units2, operation): + ul = units1.split('][') + ul[0] = ul[0][1:] + ul[-1] = ul[-1][:-1] + ol = units2.split('][') + ol[0] = ol[0][1:] + ol[-1] = ol[-1][:-1] + + usl = list() + for us in ul: + up = us.split('^') + if len(up) == 1: + up.append("1") + + if up[0]: + usl.append(tuple(up)) + + osl = list() + for os in ol: + op = os.split('^') + if len(op) == 1: + op.append("1") + + if op[0]: + osl.append(tuple(op)) + + units_elements = list() + for us in usl: + new_us = [us[0]] + i = 0 + for os in osl: + if os[0] == us[0]: + if operation == '-': + power = int(os[1]) - int(us[1]) + else: + power = int(os[1]) + int(us[1]) + del osl[i] + break + i += 1 + else: + power = int(us[1]) + + if power != 0: + new_us.append(str(power)) + units_elements.append(new_us) + + if len(osl) != 0: + units_elements += osl + + units = list() + for us in units_elements: + if us is not None: + if int(us[1]) != 1: + units.append("[" + us[0] + "^" + us[1] + "]") + else: + units.append("[" + us[0] + "]") + return "".join(units) diff --git a/qgs/params/params.py b/qgs/params/params.py index 0d33c25..5d34b41 100644 --- a/qgs/params/params.py +++ b/qgs/params/params.py @@ -46,7 +46,7 @@ import warnings from abc import ABC -from qgs.params.parameter import Parameter +from qgs.params.parameter import Parameter, ScalingParameter, ParametersArray from qgs.basis.fourier import contiguous_channel_basis, contiguous_basin_basis from qgs.basis.fourier import ChannelFourierBasis, BasinFourierBasis @@ -54,7 +54,7 @@ # TODO: - store model version in a variable somewhere -# - force the user to define the aspect ratio n at parameter object instantiation +# - force or warn the user to define the aspect ratio n at parameter object instantiation class Params(ABC): @@ -88,11 +88,23 @@ def set_params(self, dic): self.__dict__[key] = val else: d = self.__dict__[key].__dict__ - self.__dict__[key] = Parameter(val, input_dimensional=d['_input_dimensional'], + self.__dict__[key] = Parameter(val, + input_dimensional=d['_input_dimensional'], units=d['_units'], description=d['_description'], scale_object=d['_scale_object'], + symbol=d['_symbol'], return_dimensional=d['_return_dimensional']) + elif isinstance(self.__dict__[key], ScalingParameter): + if isinstance(val, ScalingParameter): + self.__dict__[key] = val + else: + d = self.__dict__[key].__dict__ + self.__dict__[key] = ScalingParameter(val, + units=d['_units'], + description=d['_description'], + symbol=d['_symbol'], + dimensional=d['_dimensional']) else: self.__dict__[key] = val @@ -113,6 +125,12 @@ def __str__(self): else: units = "[nondim]" s += "'" + key + "': " + str(efval) + " " + units + " (" + val.description + "),\n" + elif isinstance(val, ScalingParameter): + if val.dimensional: + units = val.units + else: + units = "[nondim]" + s += "'" + key + "': " + str(val) + " " + units + " (" + val.description + "),\n" elif isinstance(val, (np.ndarray, list, tuple)) and isinstance(val[0], Parameter): for i, v in enumerate(val): if v.input_dimensional: @@ -136,50 +154,6 @@ def print_params(self): """Print the parameters contained in the container.""" print(self._list_params()) - @staticmethod - def create_params_array(values, input_dimensional=None, units=None, scale_object=None, description=None, - return_dimensional=None): - - if hasattr(values, "__iter__"): - ls = len(values) - if not isinstance(input_dimensional, list): - if input_dimensional is None: - input_dimensional = True - idx = ls * [input_dimensional] - else: - idx = input_dimensional - if not isinstance(units, list): - if units is None: - units = "" - u = ls * [units] - else: - u = units - if not isinstance(description, list): - if description is None: - description = "" - d = ls * [description] - else: - d = description - if not isinstance(scale_object, list): - s = ls * [scale_object] - else: - s = scale_object - if not isinstance(return_dimensional, list): - if return_dimensional is None: - return_dimensional = False - rd = ls * [return_dimensional] - else: - rd = return_dimensional - arr = list() - for i, val in enumerate(values): - arr.append(Parameter(val, input_dimensional=idx[i], units=u[i], scale_object=s[i], description=d[i], - return_dimensional=rd[i])) - else: - arr = values * [Parameter(0.e0, input_dimensional=input_dimensional, units=units, scale_object=scale_object, - description=description, return_dimensional=return_dimensional)] - - return np.array(arr, dtype=object) - def __repr__(self): s = super(Params, self).__repr__()+"\n"+self._list_params() return s @@ -249,17 +223,16 @@ def __init__(self, dic=None): # Scale parameters for the ocean and the atmosphere # ----------------------------------------------------------- - self.scale = Parameter(5.e6, units='[m]', description="characteristic space scale (L*pi)", - return_dimensional=True) - self.f0 = Parameter(1.032e-4, units='[s^-1]', description="Coriolis parameter at the middle of the domain", - return_dimensional=True) - self.n = Parameter(1.3e0, input_dimensional=False, description="aspect ratio (n = 2 L_y / L_x)") - self.rra = Parameter(6370.e3, units='[m]', description="earth radius", return_dimensional=True) - self.phi0_npi = Parameter(0.25e0, input_dimensional=False, description="latitude expressed in fraction of pi") - self.deltap = Parameter(5.e4, units='[Pa]', description='pressure difference between the two atmospheric layers', - return_dimensional=True) - self.Ha = Parameter(8500., units='[m]', description="Average height of the 500 hPa pressure level at midlatitude", - return_dimensional=True) + self.scale = ScalingParameter(5.e6, units='[m]', description="characteristic space scale (L*pi)", dimensional=True) + self.f0 = ScalingParameter(1.032e-4, units='[s^-1]', description="Coriolis parameter at the middle of the domain", + dimensional=True, symbol=Symbol('f0')) + self.n = ScalingParameter(1.3e0, dimensional=False, description="aspect ratio (n = 2 L_y / L_x)", symbol=Symbol('n', positive=True)) + self.rra = ScalingParameter(6370.e3, units='[m]', description="earth radius", dimensional=True) + self.phi0_npi = ScalingParameter(0.25e0, dimensional=False, description="latitude expressed in fraction of pi") + self.deltap = ScalingParameter(5.e4, units='[Pa]', description='pressure difference between the two atmospheric layers', + dimensional=True) + self.Ha = ScalingParameter(8500., units='[m]', description="Average height of the 500 hPa pressure level at midlatitude", + dimensional=True, symbol=Symbol('H_a')) self.set_params(dic) # ---------------------------------------- @@ -269,35 +242,35 @@ def __init__(self, dic=None): @property def L(self): """Parameter: Typical length scale :math:`L` of the model, in meters [:math:`m`].""" - return Parameter(self.scale / np.pi, units=self.scale.units, description='Typical length scale L', - return_dimensional=True) + return ScalingParameter(self.scale / np.pi, units=self.scale.units, description='Typical length scale L', + symbol=Symbol('L'), dimensional=True) @property def L_y(self): """Parameter: The meridional extent :math:`L_y = \\pi \\, L` of the model's domain, in meters [:math:`m`].""" - return Parameter(self.scale, units=self.scale.units, description='The meridional extent of the model domain', - return_dimensional=True) + return ScalingParameter(self.scale, units=self.scale.units, description='The meridional extent of the model domain', + dimensional=True) @property def L_x(self): """Parameter: The zonal extent :math:`L_x = 2 \\pi \\, L / n` of the model's domain, in meters [:math:`m`].""" - return Parameter(2 * self.scale / self.n, units=self.scale.units, - description='The zonal extent of the model domain', - return_dimensional=True) + return ScalingParameter(2 * self.scale / self.n, units=self.scale.units, + description='The zonal extent of the model domain', + dimensional=True) @property def phi0(self): """Parameter: The reference latitude :math:`\\phi_0` at the center of the domain, expressed in radians [:math:`rad`].""" - return Parameter(self.phi0_npi * np.pi, units='[rad]', - description="The reference latitude of the center of the domain", - return_dimensional=True) + return ScalingParameter(self.phi0_npi * np.pi, units='[rad]', + description="The reference latitude of the center of the domain", + dimensional=True, symbol=Symbol('phi0')) @property def beta(self): """Parameter: The meridional gradient of the Coriolis parameter at :math:`\\phi_0`, expressed in [:math:`m^{-1} s^{-1}`]. """ return Parameter(self.L / self.rra * np.cos(self.phi0) / np.sin(self.phi0), input_dimensional=False, units='[m^-1][s^-1]', scale_object=self, - description="Meridional gradient of the Coriolis parameter at phi_0") + description="Meridional gradient of the Coriolis parameter at phi_0", symbol=Symbol('beta')) class AtmosphericParams(Params): @@ -330,11 +303,11 @@ def __init__(self, scale_params, dic=None): # Parameters for the atmosphere self.kd = Parameter(0.1, input_dimensional=False, scale_object=scale_params, units='[s^-1]', - description="atmosphere bottom friction coefficient") + description="atmosphere bottom friction coefficient", symbol=Symbol('k_d')) self.kdp = Parameter(0.01, input_dimensional=False, scale_object=scale_params, units='[s^-1]', - description="atmosphere internal friction coefficient") + description="atmosphere internal friction coefficient", symbol=Symbol('k_p')) self.sigma = Parameter(0.2e0, input_dimensional=False, scale_object=scale_params, units='[m^2][s^-2][Pa^-2]', - description="static stability of the atmosphere") + description="static stability of the atmosphere", symbol=Symbol('sigma')) self.set_params(dic) @@ -342,7 +315,7 @@ def __init__(self, scale_params, dic=None): def sig0(self): """Parameter: Static stability of the atmosphere divided by 2.""" return Parameter(self.sigma / 2, input_dimensional=False, scale_object=self._scale_params, units='[m^2][s^-2][Pa^-2]', - description="0.5 * static stability of the atmosphere") + description="0.5 * static stability of the atmosphere", symbol=self.sigma.symbol / 2) class AtmosphericTemperatureParams(Params): @@ -391,7 +364,7 @@ def __init__(self, scale_params, dic=None): self._scale_params = scale_params self.hd = Parameter(0.045, input_dimensional=False, units='[s]', scale_object=scale_params, - description="Newtonian cooling coefficient") + description="Newtonian cooling coefficient", symbol=Symbol('hd')) self.thetas = None # Radiative equilibrium mean temperature decomposition on the model's modes self.gamma = None @@ -400,10 +373,11 @@ def __init__(self, scale_params, dic=None): self.T0 = None self.sc = None self.hlambda = None + self.dynamic_T = None self.set_params(dic) - def set_insolation(self, value, pos=None): + def set_insolation(self, value, pos=None, dynamic_T=False): """Function to define the spectral decomposition of the constant short-wave radiation of the atmosphere (insolation) :math:`C_{{\\rm a}, i}` (:attr:`~.AtmosphericTemperatureParams.C`). @@ -414,33 +388,44 @@ def set_insolation(self, value, pos=None): If an iterable is provided, create a vector of spectral decomposition parameters corresponding to it. pos: int, optional Indicate in which component to set the `value`. + dynamic_T: bool, optional + Whether or not the dynamic temperature scheme is activated. """ # TODO: - check for the dimensionality of the arguments if isinstance(value, (float, int)) and pos is not None and self.C is not None: + offset = 1 + if self.dynamic_T or dynamic_T: + offset = 0 self.C[pos] = Parameter(value, units='[W][m^-2]', scale_object=self._scale_params, - description="spectral component "+str(pos+1)+" of the short-wave radiation of the atmosphere", - return_dimensional=True) + description="spectral component "+str(pos+offset)+" of the short-wave radiation of the atmosphere", + return_dimensional=True, symbol=Symbol('C_a'+str(pos+offset))) elif hasattr(value, "__iter__"): - self._create_insolation(value) + self._create_insolation(value, dynamic_T) else: warnings.warn('A scalar value was provided, but without the `pos` argument indicating in which ' + 'component of the spectral decomposition to put it: Spectral decomposition unchanged !' + 'Please specify it or give a vector as `value`.') - def _create_insolation(self, values): + def _create_insolation(self, values, dynamic_T=False): if hasattr(values, "__iter__"): dim = len(values) + values = list(values) else: dim = values values = dim * [0.] - d = ["spectral component "+str(pos+1)+" of the short-wave radiation of the atmosphere" for pos in range(dim)] + offset = 1 + if dynamic_T: + offset = 0 + self.dynamic_T = True + d = ["spectral component "+str(pos+offset)+" of the short-wave radiation of the atmosphere" for pos in range(dim)] + sy = [Symbol('C_a'+str(pos+offset)) for pos in range(dim)] - self.C = self.create_params_array(values, units='[W][m^-2]', scale_object=self._scale_params, - description=d, return_dimensional=True) + self.C = ParametersArray(values, units='[W][m^-2]', scale_object=self._scale_params, + description=d, return_dimensional=True, symbols=sy) def set_thetas(self, value, pos=None): """Function to define the spectral decomposition of the Newtonian cooling @@ -460,7 +445,7 @@ def set_thetas(self, value, pos=None): if isinstance(value, (float, int)) and pos is not None and self.thetas is not None: self.thetas[pos] = Parameter(value, scale_object=self._scale_params, description="spectral components "+str(pos+1)+" of the temperature profile", - return_dimensional=False, input_dimensional=False) + return_dimensional=False, input_dimensional=False, symbol=Symbol('thetas_'+str(pos+1))) elif hasattr(value, "__iter__"): self._create_thetas(value) else: @@ -472,14 +457,16 @@ def _create_thetas(self, values): if hasattr(values, "__iter__"): dim = len(values) + values = list(values) else: dim = values values = dim * [0.] d = ["spectral component "+str(pos+1)+" of the temperature profile" for pos in range(dim)] + sy = [Symbol('thetas_'+str(pos+1)) for pos in range(dim)] - self.thetas = self.create_params_array(values, scale_object=self._scale_params, - description=d, return_dimensional=False, input_dimensional=False) + self.thetas = ParametersArray(values, scale_object=self._scale_params, + description=d, return_dimensional=False, input_dimensional=False, symbols=sy) class OceanicParams(Params): @@ -512,13 +499,13 @@ def __init__(self, scale_params, dic=None): self._scale_params = scale_params self.gp = Parameter(3.1e-2, units='[m][s^-2]', return_dimensional=True, scale_object=scale_params, - description='reduced gravity') + description='reduced gravity', symbol=Symbol('g_p')) self.r = Parameter(1.e-8, units='[s^-1]', scale_object=scale_params, - description="frictional coefficient at the bottom of the ocean") + description="frictional coefficient at the bottom of the ocean", symbol=Symbol('r')) self.h = Parameter(5.e2, units='[m]', return_dimensional=True, scale_object=scale_params, - description="depth of the water layer of the ocean") + description="depth of the water layer of the ocean", symbol=Symbol('h')) self.d = Parameter(1.e-8, units='[s^-1]', scale_object=scale_params, - description="strength of the ocean-atmosphere mechanical coupling") + description="strength of the ocean-atmosphere mechanical coupling", symbol=Symbol('d')) self.set_params(dic) @@ -555,14 +542,15 @@ def __init__(self, scale_params, dic=None): self._scale_params = scale_params self.gamma = Parameter(2.e8, units='[J][m^-2][K^-1]', scale_object=scale_params, return_dimensional=True, - description='specific heat capacity of the ocean') + description='specific heat capacity of the ocean', symbol=Symbol('gamma_o')) self.C = None self.T0 = None + self.dynamic_T = None self.set_params(dic) - def set_insolation(self, value, pos=None): + def set_insolation(self, value, pos=None, dynamic_T=False): """Function to define the spectral decomposition of the constant short-wave radiation of the ocean (insolation) :math:`C_{{\\rm o}, i}` (:attr:`~.OceanicTemperatureParams.C`). @@ -573,31 +561,42 @@ def set_insolation(self, value, pos=None): If an iterable is provided, create a vector of spectral decomposition parameters corresponding to it. pos: int, optional Indicate in which component to set the `value`. + dynamic_T: bool, optional + Whether or not the dynamic temperature scheme is activated. """ if isinstance(value, (float, int)) and pos is not None and self.C is not None: + offset = 1 + if self.dynamic_T or dynamic_T: + offset = 0 self.C[pos] = Parameter(value, units='[W][m^-2]', scale_object=self._scale_params, - description="spectral component "+str(pos)+" of the short-wave radiation of the ocean", - return_dimensional=True) + description="spectral component "+str(pos+offset)+" of the short-wave radiation of the ocean", + return_dimensional=True, symbol=Symbol('C_go'+str(pos+offset))) elif hasattr(value, "__iter__"): - self._create_insolation(value) + self._create_insolation(value, dynamic_T) else: warnings.warn('A scalar value was provided, but without the `pos` argument indicating in which ' + 'component of the spectral decomposition to put it: Spectral decomposition unchanged !' + 'Please specify it or give a vector as `value`.') - def _create_insolation(self, values): + def _create_insolation(self, values, dynamic_T=False): if hasattr(values, "__iter__"): dim = len(values) + values = list(values) else: dim = values values = dim * [0.] - d = ["spectral component "+str(pos)+" of the short-wave radiation of the ocean" for pos in range(dim)] + offset = 1 + if dynamic_T: + offset = 0 + self.dynamic_T = True + d = ["spectral component "+str(pos+offset)+" of the short-wave radiation of the ocean" for pos in range(dim)] + sy = [Symbol('C_go'+str(pos+offset)) for pos in range(dim)] - self.C = self.create_params_array(values, units='[W][m^-2]', scale_object=self._scale_params, - description=d, return_dimensional=True) + self.C = ParametersArray(values, units='[W][m^-2]', scale_object=self._scale_params, + description=d, return_dimensional=True, symbols=sy) class GroundParams(Params): @@ -657,7 +656,7 @@ def set_orography(self, value, pos=None, basis="atmospheric"): if isinstance(value, (float, int)) and pos is not None and self.hk is not None: self.hk[pos] = Parameter(value, scale_object=self._scale_params, description="spectral components "+str(pos+1)+" of the orography", - return_dimensional=False, input_dimensional=False) + return_dimensional=False, input_dimensional=False, symbol=Symbol('hk_'+str(pos+1))) elif hasattr(value, "__iter__"): self._create_orography(value) else: @@ -669,14 +668,16 @@ def _create_orography(self, values): if hasattr(values, "__iter__"): dim = len(values) + values = list(values) else: dim = values values = dim * [0.] d = ["spectral component "+str(pos+1)+" of the orography" for pos in range(dim)] + sy = [Symbol('hk_'+str(pos+1)) for pos in range(dim)] - self.hk = self.create_params_array(values, scale_object=self._scale_params, - description=d, return_dimensional=False, input_dimensional=False) + self.hk = ParametersArray(values, scale_object=self._scale_params, + description=d, return_dimensional=False, input_dimensional=False, symbols=sy) class GroundTemperatureParams(Params): @@ -711,14 +712,15 @@ def __init__(self, scale_params, dic=None): self._scale_params = scale_params self.gamma = Parameter(2.e8, units='[J][m^-2][K^-1]', scale_object=scale_params, return_dimensional=True, - description='specific heat capacity of the ground') + description='specific heat capacity of the ground', symbol=Symbol('gamma_g')) self.C = None self.T0 = None + self.dynamic_T = None self.set_params(dic) - def set_insolation(self, value, pos=None): + def set_insolation(self, value, pos=None, dynamic_T=False): """Function to define the decomposition of the constant short-wave radiation of the ground (insolation) :math:`C_{{\\rm g}, i}` (:attr:`~.GroundTemperatureParams.C`). @@ -729,33 +731,44 @@ def set_insolation(self, value, pos=None): If an iterable is provided, create a vector of spectral decomposition parameters corresponding to it. pos: int, optional Indicate in which component to set the `value`. + dynamic_T: bool, optional + Whether or not the dynamic temperature scheme is activated. """ # TODO: - check for the dimensionality of the arguments if isinstance(value, (float, int)) and pos is not None and self.C is not None: + offset = 1 + if self.dynamic_T or dynamic_T: + offset = 0 self.C[pos] = Parameter(value, units='[W][m^-2]', scale_object=self._scale_params, - description="spectral component "+str(pos+1)+" of the short-wave radiation of the ground", - return_dimensional=True) + description="spectral component "+str(pos+offset)+" of the short-wave radiation of the ground", + return_dimensional=True, symbol=Symbol('C_go'+str(pos+offset))) elif hasattr(value, "__iter__"): - self._create_insolation(value) + self._create_insolation(value, dynamic_T) else: warnings.warn('A scalar value was provided, but without the `pos` argument indicating in which ' + 'component of the spectral decomposition to put it: Spectral decomposition unchanged !' + 'Please specify it or give a vector as `value`.') - def _create_insolation(self, values): + def _create_insolation(self, values, dynamic_T=False): if hasattr(values, "__iter__"): dim = len(values) + values = list(values) else: dim = values values = dim * [0.] - d = ["spectral component "+str(pos+1)+" of the short-wave radiation of the ground" for pos in range(dim)] + offset = 1 + if dynamic_T: + offset = 0 + self.dynamic_T = True + d = ["spectral component "+str(pos+offset)+" of the short-wave radiation of the ground" for pos in range(dim)] + sy = [Symbol('C_go'+str(pos+offset)) for pos in range(dim)] - self.C = self.create_params_array(values, units='[W][m^-2]', scale_object=self._scale_params, - description=d, return_dimensional=True) + self.C = ParametersArray(values, units='[W][m^-2]', scale_object=self._scale_params, + description=d, return_dimensional=True, symbols=sy) class QgParams(Params): @@ -919,11 +932,10 @@ def __init__(self, dic=None, scale_params=None, self.time_unit = 'days' # Physical constants - self.rr = Parameter(287.058e0, return_dimensional=True, units='[J][kg^-1][K^-1]', - scale_object=self.scale_params, description="gas constant of dry air") + scale_object=self.scale_params, description="gas constant of dry air", symbol=Symbol('R')) self.sb = Parameter(5.67e-8, return_dimensional=True, units='[J][m^-2][s^-1][K^-4]', - scale_object=self.scale_params, description="Stefan-Boltzmann constant") + scale_object=self.scale_params, description="Stefan-Boltzmann constant", symbol=Symbol('sigma_b')) self.set_params(dic) @@ -938,7 +950,7 @@ def LR(self): scp = self.scale_params if op is not None: try: - return np.sqrt(op.gp * op.h) / scp.f0 + return (op.gp * op.h) ** 0.5 / scp.f0 except: return None else: @@ -1143,10 +1155,24 @@ def set_params(self, dic): if dic is not None: for key, val in zip(dic.keys(), dic.values()): if key in self.__dict__.keys(): - self.__dict__[key] = val + if isinstance(self.__dict__[key], Parameter): + if isinstance(val, Parameter): + self.__dict__[key] = val + else: + d = self.__dict__[key].__dict__ + self.__dict__[key] = Parameter(val, + input_dimensional=d['_input_dimensional'], + units=d['_units'], + description=d['_description'], + scale_object=d['_scale_object'], + symbol=d['_symbol'], + return_dimensional=d['_return_dimensional']) + else: + self.__dict__[key] = val if 'scale_params' in self.__dict__.keys(): self.scale_params.set_params(dic) + if 'atmospheric_params' in self.__dict__.keys(): if self.atmospheric_params is not None: self.atmospheric_params.set_params(dic) @@ -1353,35 +1379,40 @@ def oceanic_basis(self, basis): self.atemperature_params.gamma = Parameter(1.e7, units='[J][m^-2][K^-1]', scale_object=self.scale_params, description='specific heat capacity of the atmosphere', - return_dimensional=True) + return_dimensional=True, symbol=Symbol('gamma_a')) if self.dynamic_T: - self.atemperature_params.set_insolation((self.nmod[0] + 1) * [0.e0]) - self.atemperature_params.set_insolation(100.0, 0) - self.atemperature_params.set_insolation(100.0, 1) + self.atemperature_params.set_insolation((self.nmod[0] + 1) * [0.e0], None, True) + self.atemperature_params.set_insolation(100.0, 0, True) + self.atemperature_params.set_insolation(100.0, 1, True) else: self.atemperature_params.set_insolation(self.nmod[0] * [0.e0]) self.atemperature_params.set_insolation(100.0, 0) self.atemperature_params.T0 = Parameter(270.0, units='[K]', scale_object=self.scale_params, return_dimensional=True, - description="stationary solution for the 0-th order atmospheric temperature") + description="stationary solution for the 0-th order atmospheric temperature", + symbol=Symbol('T_a0')) self.atemperature_params.eps = Parameter(0.76e0, input_dimensional=False, - description="emissivity coefficient for the grey-body atmosphere") + description="emissivity coefficient for the grey-body atmosphere", + symbol=Symbol('epsilon')) self.atemperature_params.sc = Parameter(1., input_dimensional=False, - description="ratio of surface to atmosphere temperature") + description="ratio of surface to atmosphere temperature", + symbol=Symbol('sc')) self.atemperature_params.hlambda = Parameter(20.00, units='[W][m^-2][K^-1]', scale_object=self.scale_params, return_dimensional=True, - description="sensible+turbulent heat exchange between ocean/ground and atmosphere") + description="sensible+turbulent heat exchange between ocean/ground and atmosphere", + symbol=Symbol('lambda')) if self.gotemperature_params is not None: if self.dynamic_T: - self.gotemperature_params.set_insolation((self.nmod[0] + 1) * [0.e0]) - self.gotemperature_params.set_insolation(350.0, 0) - self.gotemperature_params.set_insolation(350.0, 1) + self.gotemperature_params.set_insolation((self.nmod[0] + 1) * [0.e0], None, True) + self.gotemperature_params.set_insolation(350.0, 0, True) + self.gotemperature_params.set_insolation(350.0, 1, True) else: self.gotemperature_params.set_insolation(self.nmod[0] * [0.e0]) self.gotemperature_params.set_insolation(350.0, 0) self.gotemperature_params.T0 = Parameter(285.0, units='[K]', scale_object=self.scale_params, return_dimensional=True, - description="stationary solution for the 0-th order oceanic temperature") + description="stationary solution for the 0-th order oceanic temperature", + symbol=Symbol('T_go0')) # if setting an ocean, then disable the orography if self.ground_params is not None: self.ground_params.hk = None @@ -1412,24 +1443,28 @@ def ground_basis(self, basis): self.atemperature_params.gamma = Parameter(1.e7, units='[J][m^-2][K^-1]', scale_object=self.scale_params, description='specific heat capacity of the atmosphere', - return_dimensional=True) + return_dimensional=True, symbol=Symbol('gamma_a')) if self.dynamic_T: - self.atemperature_params.set_insolation((self.nmod[0] + 1) * [0.e0]) - self.atemperature_params.set_insolation(100.0, 0) - self.atemperature_params.set_insolation(100.0, 1) + self.atemperature_params.set_insolation((self.nmod[0] + 1) * [0.e0], None, True) + self.atemperature_params.set_insolation(100.0, 0, True) + self.atemperature_params.set_insolation(100.0, 1, True) else: self.atemperature_params.set_insolation(self.nmod[0] * [0.e0]) self.atemperature_params.set_insolation(100.0, 0) self.atemperature_params.T0 = Parameter(270.0, units='[K]', scale_object=self.scale_params, return_dimensional=True, - description="stationary solution for the 0-th order atmospheric temperature") + description="stationary solution for the 0-th order atmospheric temperature", + symbol=Symbol('T_a0')) self.atemperature_params.eps = Parameter(0.76e0, input_dimensional=False, - description="emissivity coefficient for the grey-body atmosphere") + description="emissivity coefficient for the grey-body atmosphere", + symbol=Symbol('epsilon')) self.atemperature_params.sc = Parameter(1., input_dimensional=False, - description="ratio of surface to atmosphere temperature") + description="ratio of surface to atmosphere temperature", + symbol=Symbol('sc')) self.atemperature_params.hlambda = Parameter(20.00, units='[W][m^-2][K^-1]', scale_object=self.scale_params, return_dimensional=True, - description="sensible+turbulent heat exchange between ocean/ground and atmosphere") + description="sensible+turbulent heat exchange between ocean/ground and atmosphere", + symbol=Symbol('lambda')) if self.gotemperature_params is not None: # if orography is disabled, enable it! @@ -1441,14 +1476,15 @@ def ground_basis(self, basis): self.ground_params.set_orography(self._number_of_ground_modes * [0.e0]) self.ground_params.set_orography(0.1, 1) if self.dynamic_T: - self.gotemperature_params.set_insolation((self.nmod[0] + 1) * [0.e0]) - self.gotemperature_params.set_insolation(350.0, 0) - self.gotemperature_params.set_insolation(350.0, 1) + self.gotemperature_params.set_insolation((self.nmod[0] + 1) * [0.e0], None, True) + self.gotemperature_params.set_insolation(350.0, 0, True) + self.gotemperature_params.set_insolation(350.0, 1, True) else: self.gotemperature_params.set_insolation(self.nmod[0] * [0.e0]) self.gotemperature_params.set_insolation(350.0, 0) self.gotemperature_params.T0 = Parameter(285.0, units='[K]', scale_object=self.scale_params, return_dimensional=True, - description="stationary solution for the 0-th order oceanic temperature") + description="stationary solution for the 0-th order oceanic temperature", + symbol=Symbol('T_go0')) def set_atmospheric_modes(self, basis, auto=False): """Function to configure the atmospheric modes (basis functions) used to project the PDEs onto. @@ -1792,19 +1828,24 @@ def oblocks(self, value): self.atemperature_params.gamma = Parameter(1.e7, units='[J][m^-2][K^-1]', scale_object=self.scale_params, description='specific heat capacity of the atmosphere', - return_dimensional=True) + return_dimensional=True, + symbol=Symbol('gamma_a')) self.atemperature_params.set_insolation(self.nmod[0] * [0.e0]) self.atemperature_params.set_insolation(100.0, 0) self.atemperature_params.eps = Parameter(0.76e0, input_dimensional=False, - description="emissivity coefficient for the grey-body atmosphere") + description="emissivity coefficient for the grey-body atmosphere", + symbol=Symbol('epsilon')) self.atemperature_params.T0 = Parameter(270.0, units='[K]', scale_object=self.scale_params, return_dimensional=True, - description="stationary solution for the 0-th order atmospheric temperature") + description="stationary solution for the 0-th order atmospheric temperature", + symbol=Symbol('T_a0')) self.atemperature_params.sc = Parameter(1., input_dimensional=False, - description="ratio of surface to atmosphere temperature") + description="ratio of surface to atmosphere temperature", + symbol=Symbol('sc')) self.atemperature_params.hlambda = Parameter(20.00, units='[W][m^-2][K^-1]', scale_object=self.scale_params, return_dimensional=True, - description="sensible+turbulent heat exchange between ocean/ground and atmosphere") + description="sensible+turbulent heat exchange between ocean/ground and atmosphere", + symbol=Symbol('lambda')) if self.gotemperature_params is not None: self._number_of_ground_modes = 0 @@ -1812,7 +1853,8 @@ def oblocks(self, value): self.gotemperature_params.set_insolation(self.nmod[0] * [0.e0]) self.gotemperature_params.set_insolation(350.0, 0) self.gotemperature_params.T0 = Parameter(285.0, units='[K]', scale_object=self.scale_params, return_dimensional=True, - description="stationary solution for the 0-th order oceanic temperature") + description="stationary solution for the 0-th order oceanic temperature", + symbol=Symbol('T_go0')) # if setting an ocean, then disable the orography if self.ground_params is not None: self.ground_params.hk = None @@ -1839,20 +1881,25 @@ def gblocks(self, value): self.atemperature_params.gamma = Parameter(1.e7, units='[J][m^-2][K^-1]', scale_object=self.scale_params, description='specific heat capacity of the atmosphere', - return_dimensional=True) + return_dimensional=True, + symbol=Symbol('gamma_a')) self.atemperature_params.set_insolation(self.nmod[0] * [0.e0]) self.atemperature_params.set_insolation(100.0, 0) self.atemperature_params.eps = Parameter(0.76e0, input_dimensional=False, - description="emissivity coefficient for the grey-body atmosphere") + description="emissivity coefficient for the grey-body atmosphere", + symbol=Symbol('epsilon')) self.atemperature_params.T0 = Parameter(270.0, units='[K]', scale_object=self.scale_params, return_dimensional=True, - description="stationary solution for the 0-th order atmospheric temperature") + description="stationary solution for the 0-th order atmospheric temperature", + symbol=Symbol('T_a0')) self.atemperature_params.sc = Parameter(1., input_dimensional=False, - description="ratio of surface to atmosphere temperature") + description="ratio of surface to atmosphere temperature", + symbol=Symbol('sc')) self.atemperature_params.hlambda = Parameter(20.00, units='[W][m^-2][K^-1]', scale_object=self.scale_params, return_dimensional=True, - description="sensible+turbulent heat exchange between ocean/ground and atmosphere") + description="sensible+turbulent heat exchange between ocean/ground and atmosphere", + symbol=Symbol('lambda')) if self.gotemperature_params is not None: gmod = 0 @@ -1872,7 +1919,8 @@ def gblocks(self, value): self.gotemperature_params.set_insolation(self.nmod[0] * [0.e0]) self.gotemperature_params.set_insolation(350.0, 0) self.gotemperature_params.T0 = Parameter(285.0, units='[K]', scale_object=self.scale_params, return_dimensional=True, - description="stationary solution for the 0-th order oceanic temperature") + description="stationary solution for the 0-th order oceanic temperature", + symbol=Symbol('T_go0')) def _set_atmospheric_analytic_fourier_modes(self, nxmax, nymax, auto=False): diff --git a/qgs/plotting/util.py b/qgs/plotting/util.py index 2eb522f..6d48620 100644 --- a/qgs/plotting/util.py +++ b/qgs/plotting/util.py @@ -21,7 +21,7 @@ def std_plot(x, mean, std, ax=None, **kwargs): std: ~numpy.ndarray 1D array of standard deviation values to represent, should have the same shape as mean. ax: None or ~matplotlib.axes.Axes - A `matplotlib`_ axes instance to plot the values. If None, create one. + A `matplotlib`_ axes instance to plot the values. If `None`, create one. kwargs: dict Keyword arguments to be given to the plot routine. diff --git a/qgs/tensors/atmo_thermo_tensor.py b/qgs/tensors/atmo_thermo_tensor.py index 1a740a7..2c9660b 100644 --- a/qgs/tensors/atmo_thermo_tensor.py +++ b/qgs/tensors/atmo_thermo_tensor.py @@ -26,13 +26,13 @@ class AtmoThermoTensor(QgsTensor): The models parameters to configure the tensor. `None` to initialize an empty tensor. Default to `None`. atmospheric_inner_products: None or AtmosphericInnerProducts, optional The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected. - If None, disable the atmospheric tendencies. Default to `None`. + If `None`, disable the atmospheric tendencies. Default to `None`. oceanic_inner_products: None or OceanicInnerProducts, optional The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected. - If None, disable the oceanic tendencies. Default to `None`. + If `None`, disable the oceanic tendencies. Default to `None`. ground_inner_products: None or GroundInnerProducts, optional The inner products of the ground basis functions on which the model's PDE ground equations are projected. - If None, disable the ground tendencies. Default to `None`. + If `None`, disable the ground tendencies. Default to `None`. Attributes ---------- @@ -40,13 +40,13 @@ class AtmoThermoTensor(QgsTensor): The models parameters used to configure the tensor. `None` for an empty tensor. atmospheric_inner_products: None or AtmosphericInnerProducts The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected. - If None, disable the atmospheric tendencies. Default to `None`. + If `None`, disable the atmospheric tendencies. Default to `None`. oceanic_inner_products: None or OceanicInnerProducts The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected. - If None, disable the oceanic tendencies. Default to `None`. + If `None`, disable the oceanic tendencies. Default to `None`. ground_inner_products: None or GroundInnerProducts The inner products of the ground basis functions on which the model's PDE ground equations are projected. - If None, disable the ground tendencies. Default to `None`. + If `None`, disable the ground tendencies. Default to `None`. tensor: sparse.COO(float) The tensor :math:`\\mathcal{T}_{i,j,k}` :math:`i`-th components. jacobian_tensor: sparse.COO(float) @@ -233,13 +233,13 @@ class AtmoThermoTensorDynamicT(AtmoThermoTensor): The models parameters to configure the tensor. `None` to initialize an empty tensor. Default to `None`. atmospheric_inner_products: None or AtmosphericInnerProducts, optional The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected. - If None, disable the atmospheric tendencies. Default to `None`. + If `None`, disable the atmospheric tendencies. Default to `None`. oceanic_inner_products: None or OceanicInnerProducts, optional The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected. - If None, disable the oceanic tendencies. Default to `None`. + If `None`, disable the oceanic tendencies. Default to `None`. ground_inner_products: None or GroundInnerProducts, optional The inner products of the ground basis functions on which the model's PDE ground equations are projected. - If None, disable the ground tendencies. Default to `None`. + If `None`, disable the ground tendencies. Default to `None`. Attributes ---------- @@ -247,13 +247,13 @@ class AtmoThermoTensorDynamicT(AtmoThermoTensor): The models parameters used to configure the tensor. `None` for an empty tensor. atmospheric_inner_products: None or AtmosphericInnerProducts The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected. - If None, disable the atmospheric tendencies. Default to `None`. + If `None`, disable the atmospheric tendencies. Default to `None`. oceanic_inner_products: None or OceanicInnerProducts The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected. - If None, disable the oceanic tendencies. Default to `None`. + If `None`, disable the oceanic tendencies. Default to `None`. ground_inner_products: None or GroundInnerProducts The inner products of the ground basis functions on which the model's PDE ground equations are projected. - If None, disable the ground tendencies. Default to `None`. + If `None`, disable the ground tendencies. Default to `None`. tensor: sparse.COO(float) The tensor :math:`\\mathcal{T}_{i,j,k}` :math:`i`-th components. jacobian_tensor: sparse.COO(float) @@ -455,13 +455,13 @@ class AtmoThermoTensorT4(AtmoThermoTensorDynamicT): The models parameters to configure the tensor. `None` to initialize an empty tensor. Default to `None`. atmospheric_inner_products: None or AtmosphericInnerProducts, optional The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected. - If None, disable the atmospheric tendencies. Default to `None`. + If `None`, disable the atmospheric tendencies. Default to `None`. oceanic_inner_products: None or OceanicInnerProducts, optional The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected. - If None, disable the oceanic tendencies. Default to `None`. + If `None`, disable the oceanic tendencies. Default to `None`. ground_inner_products: None or GroundInnerProducts, optional The inner products of the ground basis functions on which the model's PDE ground equations are projected. - If None, disable the ground tendencies. Default to `None`. + If `None`, disable the ground tendencies. Default to `None`. Attributes ---------- @@ -469,13 +469,13 @@ class AtmoThermoTensorT4(AtmoThermoTensorDynamicT): The models parameters used to configure the tensor. `None` for an empty tensor. atmospheric_inner_products: None or AtmosphericInnerProducts The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected. - If None, disable the atmospheric tendencies. Default to `None`. + If `None`, disable the atmospheric tendencies. Default to `None`. oceanic_inner_products: None or OceanicInnerProducts The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected. - If None, disable the oceanic tendencies. Default to `None`. + If `None`, disable the oceanic tendencies. Default to `None`. ground_inner_products: None or GroundInnerProducts The inner products of the ground basis functions on which the model's PDE ground equations are projected. - If None, disable the ground tendencies. Default to `None`. + If `None`, disable the ground tendencies. Default to `None`. tensor: sparse.COO(float) The tensor :math:`\\mathcal{T}_{i,j,k}` :math:`i`-th components. jacobian_tensor: sparse.COO(float) diff --git a/qgs/tensors/qgtensor.py b/qgs/tensors/qgtensor.py index 7590743..5ecfdd7 100644 --- a/qgs/tensors/qgtensor.py +++ b/qgs/tensors/qgtensor.py @@ -25,13 +25,13 @@ class QgsTensor(object): The models parameters to configure the tensor. `None` to initialize an empty tensor. Default to `None`. atmospheric_inner_products: None or AtmosphericInnerProducts, optional The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected. - If None, disable the atmospheric tendencies. Default to `None`. + If `None`, disable the atmospheric tendencies. Default to `None`. oceanic_inner_products: None or OceanicInnerProducts, optional The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected. - If None, disable the oceanic tendencies. Default to `None`. + If `None`, disable the oceanic tendencies. Default to `None`. ground_inner_products: None or GroundInnerProducts, optional The inner products of the ground basis functions on which the model's PDE ground equations are projected. - If None, disable the ground tendencies. Default to `None`. + If `None`, disable the ground tendencies. Default to `None`. Attributes ---------- @@ -39,13 +39,13 @@ class QgsTensor(object): The models parameters used to configure the tensor. `None` for an empty tensor. atmospheric_inner_products: None or AtmosphericInnerProducts The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected. - If None, disable the atmospheric tendencies. Default to `None`. + If `None`, disable the atmospheric tendencies. Default to `None`. oceanic_inner_products: None or OceanicInnerProducts The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected. - If None, disable the oceanic tendencies. Default to `None`. + If `None`, disable the oceanic tendencies. Default to `None`. ground_inner_products: None or GroundInnerProducts The inner products of the ground basis functions on which the model's PDE ground equations are projected. - If None, disable the ground tendencies. Default to `None`. + If `None`, disable the ground tendencies. Default to `None`. tensor: sparse.COO(float) The tensor :math:`\\mathcal{T}_{i,j,k}` :math:`i`-th components. jacobian_tensor: sparse.COO(float) @@ -703,12 +703,12 @@ def jacobian_from_tensor(tensor): Parameters ---------- - tensor: ~sparse.COO + tensor: sparse.COO The qgs tensor. Returns ------- - ~sparse.COO + sparse.COO The Jacobian tensor. """ @@ -729,12 +729,12 @@ def simplify_tensor(tensor): Parameters ---------- - tensor: ~sparse.COO + tensor: sparse.COO The tensor to simplify. Returns ------- - ~sparse.COO + sparse.COO The upper-triangularized tensor. """ coords = tensor.coords.copy() @@ -849,13 +849,13 @@ class QgsTensorDynamicT(QgsTensor): The models parameters to configure the tensor. `None` to initialize an empty tensor. Default to `None`. atmospheric_inner_products: None or AtmosphericInnerProducts, optional The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected. - If None, disable the atmospheric tendencies. Default to `None`. + If `None`, disable the atmospheric tendencies. Default to `None`. oceanic_inner_products: None or OceanicInnerProducts, optional The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected. - If None, disable the oceanic tendencies. Default to `None`. + If `None`, disable the oceanic tendencies. Default to `None`. ground_inner_products: None or GroundInnerProducts, optional The inner products of the ground basis functions on which the model's PDE ground equations are projected. - If None, disable the ground tendencies. Default to `None`. + If `None`, disable the ground tendencies. Default to `None`. Attributes ---------- @@ -863,13 +863,13 @@ class QgsTensorDynamicT(QgsTensor): The models parameters used to configure the tensor. `None` for an empty tensor. atmospheric_inner_products: None or AtmosphericInnerProducts The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected. - If None, disable the atmospheric tendencies. Default to `None`. + If `None`, disable the atmospheric tendencies. Default to `None`. oceanic_inner_products: None or OceanicInnerProducts The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected. - If None, disable the oceanic tendencies. Default to `None`. + If `None`, disable the oceanic tendencies. Default to `None`. ground_inner_products: None or GroundInnerProducts The inner products of the ground basis functions on which the model's PDE ground equations are projected. - If None, disable the ground tendencies. Default to `None`. + If `None`, disable the ground tendencies. Default to `None`. tensor: sparse.COO(float) The tensor :math:`\\mathcal{T}_{i,j,k}` :math:`i`-th components. jacobian_tensor: sparse.COO(float) @@ -977,7 +977,6 @@ def _compute_stored_full_dict(self): sparse_arrays_full_dict[self._theta_a(i)].append(self._shift_tensor_coordinates(- par.T4LSBpgo * val, self._deltaT_o(0))) if ground_temp: - if par.T4LSBpgo is not None: val = sp.tensordot(a_theta[i], aips._v, axes=1) if val.nnz > 0: @@ -1180,13 +1179,13 @@ class QgsTensorT4(QgsTensorDynamicT): The models parameters to configure the tensor. `None` to initialize an empty tensor. Default to `None`. atmospheric_inner_products: None or AtmosphericInnerProducts, optional The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected. - If None, disable the atmospheric tendencies. Default to `None`. + If `None`, disable the atmospheric tendencies. Default to `None`. oceanic_inner_products: None or OceanicInnerProducts, optional The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected. - If None, disable the oceanic tendencies. Default to `None`. + If `None`, disable the oceanic tendencies. Default to `None`. ground_inner_products: None or GroundInnerProducts, optional The inner products of the ground basis functions on which the model's PDE ground equations are projected. - If None, disable the ground tendencies. Default to `None`. + If `None`, disable the ground tendencies. Default to `None`. Attributes ---------- @@ -1194,13 +1193,13 @@ class QgsTensorT4(QgsTensorDynamicT): The models parameters used to configure the tensor. `None` for an empty tensor. atmospheric_inner_products: None or AtmosphericInnerProducts The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected. - If None, disable the atmospheric tendencies. Default to `None`. + If `None`, disable the atmospheric tendencies. Default to `None`. oceanic_inner_products: None or OceanicInnerProducts The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected. - If None, disable the oceanic tendencies. Default to `None`. + If `None`, disable the oceanic tendencies. Default to `None`. ground_inner_products: None or GroundInnerProducts The inner products of the ground basis functions on which the model's PDE ground equations are projected. - If None, disable the ground tendencies. Default to `None`. + If `None`, disable the ground tendencies. Default to `None`. tensor: sparse.COO(float) The tensor :math:`\\mathcal{T}_{i,j,k}` :math:`i`-th components. jacobian_tensor: sparse.COO(float) diff --git a/qgs/tensors/symbolic_qgtensor.py b/qgs/tensors/symbolic_qgtensor.py new file mode 100644 index 0000000..64b0111 --- /dev/null +++ b/qgs/tensors/symbolic_qgtensor.py @@ -0,0 +1,1529 @@ +""" + symbolic qgs tensor module + ========================== + + This module computes and holds the symbolic representation of the tensors representing the tendencies of the model's + equations. + +""" +from qgs.functions.symbolic_mul import add_to_dict, symbolic_tensordot +from qgs.params.params import Parameter, ScalingParameter, ParametersArray, Params + +import numpy as np +from sympy import simplify +import pickle + +from sympy.matrices.immutable import ImmutableSparseMatrix +from sympy.tensor.array import ImmutableSparseNDimArray + +# TODO: Check non stored IP version of this + + +class SymbolicQgsTensor(object): + """Symbolic qgs tendencies tensor class. + + Parameters + ---------- + params: None or QgParams, optional + The models parameters to configure the tensor. `None` to initialize an empty tensor. Default to `None`. + atmospheric_inner_products: None or AtmosphericInnerProducts, optional + The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are + projected. + If `None`, disable the atmospheric tendencies. Default to `None`. + The inner product is returned in symbolic or numeric form. + oceanic_inner_products: None or OceanicInnerProducts, optional + The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected. + If `None`, disable the oceanic tendencies. Default to `None`. + The inner product is returned in symbolic or numeric form. + ground_inner_products: None or GroundInnerProducts, optional + The inner products of the ground basis functions on which the model's PDE ground equations are projected. + If `None`, disable the ground tendencies. Default to `None`. + The inner product is returned in symbolic or numeric form. + + Attributes + ---------- + params: None or QgParams + The models parameters used to configure the tensor. `None` for an empty tensor. + atmospheric_inner_products: None or AtmosphericInnerProducts + The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are + projected. + If `None`, disable the atmospheric tendencies. Default to `None`. + oceanic_inner_products: None or OceanicInnerProducts + The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected. + If `None`, disable the oceanic tendencies. Default to `None`. + ground_inner_products: None or GroundInnerProducts + The inner products of the ground basis functions on which the model's PDE ground equations are projected. + If `None`, disable the ground tendencies. Default to `None`. + tensor: sparse.COO(float) + The tensor :math:`\\mathcal{T}_{i,j,k}` :math:`i`-th components. + jacobian_tensor: sparse.COO(float) + The jacobian tensor :math:`\\mathcal{T}_{i,j,k} + \\mathcal{T}_{i,k,j}` :math:`i`-th components. + """ + + def __init__(self, params=None, atmospheric_inner_products=None, oceanic_inner_products=None, + ground_inner_products=None): + + self.atmospheric_inner_products = atmospheric_inner_products + self.oceanic_inner_products = oceanic_inner_products + self.ground_inner_products = ground_inner_products + self.params = params + + self.tensor = None + self.jacobian_tensor = None + + if not self.params.dynamic_T: + self.compute_tensor() + + def _psi_a(self, i): + """Transform the :math:`\\psi_{\\mathrm a}` :math:`i`-th coefficient into the effective model's variable. + + Parameters + ---------- + i: int + The :math:`i`-th coefficients of :math:`\\psi_{\\mathrm a}` + + Returns + ------- + int + The effective model's variable. + """ + + return i + 1 + + def _theta_a(self, i): + """Transform the :math:`\\theta_{\\mathrm a}` :math:`i`-th coefficient into the effective model's variable. + + Parameters + ---------- + i: int + The :math:`i`-th coefficients of :math:`\\theta_{\\mathrm a}` + + Returns + ------- + int + The effective model's variable. + """ + return i + self.params.variables_range[0] + 1 + + def _psi_o(self, i): + """Transform the :math:`\\psi_{\\mathrm o}` :math:`i`-th coefficient into the effective model's variable. + + Parameters + ---------- + i: int + The :math:`i`-th coefficients of :math:`\\psi_{\\mathrm o}` + + Returns + ------- + int + The effective model's variable. + """ + return i + self.params.variables_range[1] + 1 + + def _deltaT_o(self, i): + """Transform the :math:`\\delta T_{\\mathrm o}` :math:`i`-th coefficient into the effective model's variable. + + Parameters + ---------- + i: int + The :math:`i`-th coefficients of :math:`\\delta T_{\\mathrm o}` + + Returns + ------- + int + The effective model's variable. + """ + return i + self.params.variables_range[2] + 1 + + def _deltaT_g(self, i): + """Transform the :math:`\\delta T_{\\mathrm o}` :math:`i`-th coefficient into the effective model's variable. + + Parameters + ---------- + i: int + The :math:`i`-th coefficients of :math:`\\delta T_{\\mathrm o}` + + Returns + ------- + int + The effective model's variable. + """ + return i + self.params.variables_range[1] + 1 + + def _compute_tensor_dicts(self): + if self.params is None: + return None + + if self.atmospheric_inner_products is None and self.oceanic_inner_products is None \ + and self.ground_inner_products is None: + return None + + aips = self.atmospheric_inner_products + + par = self.params + atp = par.atemperature_params + gp = par.ground_params + ap = par.atmospheric_params + + nvar = par.number_of_variables + + bips = None + if self.oceanic_inner_products is not None: + bips = self.oceanic_inner_products + ocean = True + else: + ocean = False + + if self.ground_inner_products is not None: + bips = self.ground_inner_products + ground_temp = True + else: + ground_temp = False + + if self.params.dynamic_T: + offset = 1 + else: + offset = 0 + + # constructing some derived matrices + if aips is not None: + a_inv = dict() + for i in range(offset, nvar[1]): + for j in range(offset, nvar[1]): + a_inv[(i - offset, j - offset)] = aips.a(i, j) + + a_inv = ImmutableSparseMatrix(nvar[0], nvar[0], a_inv) + a_inv = a_inv.inverse() + a_inv = a_inv.simplify() + + a_theta = dict() + for i in range(nvar[1]): + for j in range(nvar[1]): + a_theta[(i, j)] = ap.sig0.symbol * aips.a(i, j) - aips.u(i, j) + + a_theta = ImmutableSparseMatrix(nvar[1], nvar[1], a_theta) + a_theta = a_theta.inverse() + a_theta = a_theta.simplify() + + if bips is not None: + if ocean: + U_inv = dict() + for i in range(nvar[3]): + for j in range(nvar[3]): + U_inv[(i, j)] = bips.U(i, j) + U_inv = ImmutableSparseMatrix(nvar[3], nvar[3], U_inv) + U_inv = U_inv.inverse() + U_inv = U_inv.simplify() + + M_psio = dict() + for i in range(offset, nvar[3]): + for j in range(offset, nvar[3]): + M_psio[(i - offset, j - offset)] = bips.M(i, j) + self.params.G.symbolic_expression * bips.U(i, j) + + M_psio = ImmutableSparseMatrix(nvar[2], nvar[2], M_psio) + M_psio = M_psio.inverse() + M_psio = M_psio.simplify() + + else: + U_inv = dict() + for i in range(nvar[2]): + for j in range(nvar[2]): + U_inv[(i, j)] = bips.U(i, j) + U_inv = ImmutableSparseMatrix(nvar[2], nvar[2], U_inv) + U_inv = U_inv.inverse() + U_inv = U_inv.simplify() + + ################ + + if bips is not None: + go = bips.stored + else: + go = True + + sy_arr_dic = dict() + + if aips.stored and go: + # psi_a part + a_inv_mult_c = a_inv @ aips._c[offset:, offset:] + + if gp is not None: + hk_sym_arr = ImmutableSparseMatrix(gp.hk.symbols) + + if gp.orographic_basis == "atmospheric": + a_inv_mult_g = symbolic_tensordot(a_inv, aips._g[offset:, offset:, offset:], axes=1) + oro = symbolic_tensordot(a_inv_mult_g, hk_sym_arr, axes=1)[:, :, 0] + else: + a_inv_mult_gh = symbolic_tensordot(a_inv, aips._gh[offset:, offset:, offset:], axes=1) + oro = symbolic_tensordot(a_inv_mult_gh, hk_sym_arr, axes=1)[:, :, 0] + + a_inv_mult_b = symbolic_tensordot(a_inv, aips._b[offset:, offset:, offset:], axes=1) + + if ocean: + a_inv_mult_d = a_inv @ aips._d[offset:, offset:] + + for i in range(nvar[0]): + for j in range(nvar[0]): + jo = j + offset # skipping the theta 0 variable if it exists + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_a(i), self._psi_a(j), 0), -a_inv_mult_c[i, j] * par.scale_params.beta.symbol) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_a(i), self._psi_a(j), 0), -(ap.kd.symbol * _kronecker_delta(i, j)) / 2) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_a(i), self._theta_a(jo), 0), (ap.kd.symbol * _kronecker_delta(i, j)) / 2) + + if gp is not None: + # convert + if gp.hk is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_a(i), self._psi_a(j), 0), -oro[i, j] / 2) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_a(i), self._theta_a(jo), 0), oro[i, j] / 2) + + for k in range(nvar[0]): + ko = k + offset # skipping the theta 0 variable if it exists + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_a(i), self._psi_a(j), self._psi_a(k)), -a_inv_mult_b[i, j, k]) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_a(i), self._theta_a(jo), self._theta_a(ko)), -a_inv_mult_b[i, j, k]) + + if ocean: + for j in range(nvar[2]): + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_a(i), self._psi_o(j), 0), a_inv_mult_d[i, j] * ap.kd.symbol / 2) + + # theta_a part + a_theta_mult_u = a_theta @ aips._u + if self.params.Cpa is not None: + val_Cpa = a_theta_mult_u @ ImmutableSparseMatrix(self.params.Cpa.symbolic_expressions) + + if atp.hd is not None and atp.thetas is not None: + thetas_sym_arr = ImmutableSparseMatrix(atp.thetas.symbols) + val_thetas = a_theta_mult_u @ thetas_sym_arr + + a_theta_mult_a = a_theta @ aips._a[:, offset:] + a_theta_mult_c = a_theta @ aips._c[:, offset:] + + a_theta_mult_g = symbolic_tensordot(a_theta, aips._g[:, offset:, offset:], axes=1) + + if gp is not None: + if gp.orographic_basis == "atmospheric": + oro = symbolic_tensordot(a_theta_mult_g, hk_sym_arr, axes=1)[:, :, 0] + else: + a_theta_mult_gh = symbolic_tensordot(a_theta, aips._gh[:, offset:, offset:], axes=1) + oro = symbolic_tensordot(a_theta_mult_gh, hk_sym_arr, axes=1)[:, :, 0] + + a_theta_mult_b = symbolic_tensordot(a_theta, aips._b[:, offset:, offset:], axes=1) + + if ocean: + a_theta_mult_d = a_theta @ aips._d[:, offset:] + a_theta_mult_s = a_theta @ aips._s + + if ground_temp: + a_theta_mult_s = a_theta @ aips._s + + for i in range(nvar[1]): + if self.params.Cpa is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), 0, 0), -val_Cpa[i]) + + if atp.hd is not None and atp.thetas is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), 0, 0), -val_thetas[i] * atp.hd.symbol) + + for j in range(nvar[0]): + + jo = j + offset # skipping the theta 0 variable if it exists + + val_2 = a_theta_mult_a[i, j] * ap.kd.symbol * ap.sig0.symbol / 2 + val_3 = a_theta_mult_a[i, j] * (ap.kd.symbol / 2 + 2 * ap.kdp.symbol) * ap.sig0.symbol + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._psi_a(j), 0), val_2) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._theta_a(jo), 0), -val_3) + + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._theta_a(jo), 0), -a_theta_mult_c[i, j] * par.scale_params.beta.symbol * ap.sig0.symbol) + + if gp is not None: + if gp.hk is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._theta_a(jo), 0), -ap.sig0.symbol * oro[i, j] / 2) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._psi_a(j), 0), ap.sig0.symbol * oro[i, j] / 2) + + for k in range(nvar[0]): + ko = k + offset # skipping the theta 0 variable if it exists + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._psi_a(j), self._theta_a(ko)), - a_theta_mult_b[i, j, k] * ap.sig0.symbol) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._theta_a(jo), self._psi_a(k)), - a_theta_mult_b[i, j, k] * ap.sig0.symbol) + + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._psi_a(j), self._theta_a(ko)), a_theta_mult_g[i, j, k]) + + for j in range(nvar[1]): + if self.params.Lpa is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._theta_a(j), 0), a_theta_mult_u[i, j] * atp.sc.symbol * self.params.Lpa.symbolic_expression) + if self.params.LSBpa is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._theta_a(j), 0), a_theta_mult_u[i, j] * self.params.LSBpa.symbolic_expression) + + if atp.hd is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._theta_a(j), 0), a_theta_mult_u[i, j] * atp.hd) + + if ocean: + for j in range(nvar[2]): + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._psi_o(j), 0), -a_theta_mult_d[i, j] * ap.sig0.symbol * ap.kd.symbol / 2) + + if self.params.Lpa is not None: + for j in range(nvar[3]): + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._deltaT_o(j), 0), -a_theta_mult_s[i, j] * self.params.Lpa.symbolic_expression / 2) + if self.params.LSBpgo is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._deltaT_o(j), 0), -a_theta_mult_s[i, j] * self.params.LSBpgo.symbolic_expression) + + if ground_temp: + if self.params.Lpa is not None: + for j in range(nvar[2]): + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._deltaT_g(j), 0), -a_theta_mult_s[i, j] * self.params.Lpa.symbolic_expression / 2) + if self.params.LSBpgo is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._deltaT_g(j), 0), -a_theta_mult_s[i, j] * self.params.LSBpgo.symbolic_expression) + + if ocean: + # psi_o part + M_psio_mult_K = M_psio @ bips._K[offset:, offset:] + M_psio_mult_N = M_psio @ bips._N[offset:, offset:] + M_psio_mult_M = M_psio @ bips._M[offset:, offset:] + M_psio_mult_C = symbolic_tensordot(M_psio, bips._C[offset:, offset:, offset:], axes=1) + + for i in range(nvar[2]): + for j in range(nvar[0]): + jo = j + offset # skipping the theta 0 variable if it exists + + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_o(i), self._psi_a(j), 0), M_psio_mult_K[i, j] * par.oceanic_params.d.symbol) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_o(i), self._theta_a(jo), 0), -M_psio_mult_K[i, j] * par.oceanic_params.d.symbol) + + for j in range(nvar[2]): + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_o(i), self._psi_o(j), 0), -M_psio_mult_N[i, j] * par.scale_params.beta.symbol) + + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_o(i), self._psi_o(j), 0), -M_psio_mult_M[i, j] * (par.oceanic_params.r.symbol + par.oceanic_params.d.symbol)) + + for k in range(nvar[2]): + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_o(i), self._psi_o(j), self._psi_o(k)), - M_psio_mult_C[i, j, k]) + + # deltaT_o part + U_inv_mult_W = U_inv @ bips._W + Cpgo_sym_arr = ImmutableSparseMatrix(self.params.Cpgo.symbolic_expressions) + U_inv_mult_W_Cpgo = U_inv_mult_W @ Cpgo_sym_arr + + U_inv_mult_O = symbolic_tensordot(U_inv, bips._O[:, offset:, offset:], axes=1) + + for i in range(nvar[3]): + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_o(i), 0, 0), U_inv_mult_W_Cpgo[i]) + + for j in range(nvar[1]): + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_o(i), self._theta_a(j), 0), U_inv_mult_W[i, j] * 2 * atp.sc.symbol * self.params.Lpgo.symbolic_expression) + if self.params.sbpa is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_o(i), self._theta_a(j), 0), U_inv_mult_W[i, j] * self.params.sbpa.symbolic_expression) + + for j in range(nvar[3]): + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_o(i), self._deltaT_o(j), 0), - self.params.Lpgo.symbolic_expression * _kronecker_delta(i, j)) + if self.params.sbpgo is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_o(i), self._deltaT_o(j), 0), - self.params.sbpgo.symbolic_expression * _kronecker_delta(i, j)) + + for j in range(nvar[2]): + for k in range(nvar[2]): + ko = k + offset # skipping the T 0 variable if it exists + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_o(i), self._psi_o(j), self._deltaT_o(ko)), -U_inv_mult_O[i, j, k]) + + # deltaT_g part + if ground_temp: + U_inv_mult_W = U_inv @ bips._W + Cpgo_sym_arr = ImmutableSparseMatrix(self.params.Cpgo.symbolic_expressions) + U_inv_mult_W_Cpgo = U_inv_mult_W @ Cpgo_sym_arr + for i in range(nvar[2]): + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_g(i), 0, 0), U_inv_mult_W_Cpgo[i]) + + for j in range(nvar[1]): + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_g(i), self._theta_a(j), 0), U_inv_mult_W[i, j] * 2 * atp.sc.symbol * self.params.Lpgo.symbolic_expression) + if self.params.sbpa is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_g(i), self._theta_a(j), 0), U_inv_mult_W[i, j] * self.params.sbpa.symbolic_expression) + + for j in range(nvar[2]): + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_g(i), self._deltaT_g(j), 0), - self.params.Lpgo.symbolic_expression * _kronecker_delta(i, j)) + if self.params.sbpgo is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_g(i), self._deltaT_g(j), 0), - self.params.sbpgo.symbolic_expression * _kronecker_delta(i, j)) + + else: + # psi_a part + for i in range(nvar[0]): + for j in range(nvar[0]): + + jo = j + offset + + val = 0 + for jj in range(nvar[0]): + val += a_inv[i, jj] * aips.c(offset + jj, jo) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_a(i), self._psi_a(j), 0), - val * par.scale_params.beta.symbol) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_a(i), self._psi_a(j), 0), - (ap.kd.symbol * _kronecker_delta(i, j)) / 2) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_a(i), self._theta_a(jo), 0), (ap.kd.symbol * _kronecker_delta(i, j)) / 2) + + if gp is not None: + hk_sym_arr = ImmutableSparseNDimArray(gp.hk.symbols) + if gp.hk is not None: + oro = 0 + if gp.orographic_basis == "atmospheric": + for jj in range(nvar[0]): + for kk in range(nvar[0]): + oro += a_inv[i, jj] * aips.g(offset + jj, j, offset + kk) * hk_sym_arr[kk] + else: + for jj in range(nvar[0]): + for kk in range(nvar[0]): + oro += a_inv[i, jj] * aips.gh(offset + jj, j, offset + kk) * hk_sym_arr[kk] + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_a(i), self._psi_a(j), 0), - oro / 2) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_a(i), self._theta_a(jo), 0), oro / 2) + + for k in range(nvar[0]): + ko = k + offset # skipping the theta 0 variable if it exists + val = 0 + for jj in range(nvar[0]): + val += a_inv[i, jj] * aips.b(offset + jj, jo, ko) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_a(i), self._psi_a(j), self._psi_a(k)), - val) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_a(i), self._theta_a(jo), self._theta_a(ko)), - val) + if ocean: + for j in range(nvar[2]): + jo = j + offset + val = 0 + for jj in range(nvar[0]): + val += a_inv[i, jj] * aips.d(offset + jj, jo) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_a(i), self._psi_o(j), 0), val * ap.kd.symbol / 2) + + # theta_a part + for i in range(nvar[1]): + if self.params.Cpa is not None: + val = 0 + for jj in range(nvar[1]): + for kk in range(nvar[1]): + val -= a_theta[i, jj] * aips.u(jj, kk) * self.params.Cpa.symbolic_expressions[kk] + + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), 0, 0), val) + + if atp.hd is not None and atp.thetas is not None: + thetas_sym_arr = ImmutableSparseNDimArray(atp.thetas.symbols) + val = 0 + for jj in range(nvar[1]): + for kk in range(nvar[1]): + val -= a_theta[i, jj] * aips.u(jj, kk) * thetas_sym_arr[kk] + + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), 0, 0), val * atp.hd.symbol) + + for j in range(nvar[0]): + + jo = j + offset # skipping the theta 0 variable if it exists + + val = 0 + for jj in range(nvar[1]): + val += a_theta[i, jj] * aips.a(jj, jo) + + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._psi_a(j), 0), val * ap.kd.symbol * ap.sig0.symbol / 2) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._theta_a(jo), 0), - val * (ap.kd.symbol / 2 - 2 * ap.kdp.symbol) * ap.sig0.symbol) + + val = 0 + for jj in range(nvar[1]): + val -= a_theta[i, jj] * aips.c(jj, jo) + + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._theta_a(jo), 0), val * par.scale_params.beta.symbol * ap.sig0.symbol) + + if gp is not None: + if gp.hk is not None: + oro = 0 + if gp.orographic_basis == "atmospheric": + for jj in range(nvar[1]): + for kk in range(nvar[0]): + oro += a_theta[i, jj] * aips.g(jj, jo, offset + kk) * gp.hk[kk].symbol + else: + for jj in range(nvar[1]): + for kk in range(nvar[0]): + oro += a_theta[i, jj] * aips.gh(jj, jo, offset + kk) * gp.hk[kk].symbol + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._theta_a(jo), 0), - ap.sig0.symbol * oro / 2) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._psi_a(j), 0), ap.sig0.symbol * oro / 2) + + for k in range(nvar[0]): + ko = k + offset # skipping the theta 0 variable if it exists + val = 0 + for jj in range(nvar[1]): + val += a_theta[i, jj] * aips.b(jj, jo, ko) + + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._psi_a(j), self._theta_a(ko)), - val * ap.sig0.symbol) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._theta_a(jo), self._psi_a(k)), - val * ap.sig0.symbol) + + val = 0 + for jj in range(nvar[1]): + val += a_theta[i, jj] * aips.g(jj, jo, ko) + + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._psi_a(j), self._theta_a(ko)), val) + + for j in range(nvar[1]): + val = 0 + for jj in range(nvar[1]): + val += a_theta[i, jj] * aips.u(jj, j) + + if self.params.Lpa is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._theta_a(j), 0), val * atp.sc.symbol * self.params.Lpa.symbolic_expression) + + if self.params.LSBpa is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._theta_a(j), 0), val * self.params.LSBpa.symbolic_expression) + + if atp.hd is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._theta_a(j), 0), val * atp.hd.symbol) + + if ocean: + for j in range(nvar[2]): + jo = j + offset # skipping the theta 0 variable if it exists + val = 0 + for jj in range(nvar[1]): + val -= a_theta[i, jj] * aips.d(jj, jo) + + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._psi_o(j), 0), val * ap.sig0.symbol * ap.kd.symbol / 2) + + if self.params.Lpa is not None: + for j in range(nvar[3]): + val = 0 + for jj in range(nvar[1]): + val -= a_theta[i, jj] * aips.s(jj, j) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._deltaT_o(j), 0), val * self.params.Lpa.symbolic_expression / 2) + if self.params.LSBpgo is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._deltaT_o(j), 0), val * self.params.LSBpgo.symbolic_expression) + + if ground_temp: + if self.params.Lpa is not None: + for j in range(nvar[2]): + val = 0 + for jj in range(nvar[1]): + val -= a_theta[i, jj] * aips.s(jj, j) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._deltaT_g(j), 0), val * self.params.Lpa.symbolic_expression / 2) + if self.params.LSBpgo is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._deltaT_g(j), 0), val * self.params.LSBpgo.symbolic_expression) + + if ocean: + # psi_o part + for i in range(nvar[2]): + for j in range(nvar[0]): + jo = j + offset # skipping the theta 0 variable if it exists + + for jj in range(nvar[2]): + val = M_psio[i, jj] * bips.K(offset + jj, jo) * par.oceanic_params.d.symbol + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_o(i), self._psi_a(j), 0), val) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_o(i), self._theta_a(jo), 0), - val) + + for j in range(nvar[2]): + jo = j + offset # skipping the T 0 variable if it exists + + val = 0 + for jj in range(nvar[2]): + val -= M_psio[i, jj] * bips.N(offset + jj, jo) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_o(i), self._psi_o(j), 0), val * par.scale_params.beta.symbol) + + val = 0 + for jj in range(nvar[2]): + val -= M_psio[i, jj] * bips.M(offset + jj, jo) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_o(i), self._psi_o(j), 0), val * (par.oceanic_params.r.symbol + par.oceanic_params.d.symbol)) + + for k in range(nvar[2]): + ko = k + offset # skipping the T 0 variable if it exists + val = 0 + for jj in range(nvar[2]): + val -= M_psio[i, jj] * bips.C(offset + jj, jo, ko) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._psi_o(i), self._psi_o(j), self._psi_o(k)), val) + + # deltaT_o part + for i in range(nvar[3]): + val = 0 + for jj in range(nvar[1]): + for kk in range(nvar[3]): + val += U_inv[i, kk] * bips.W(kk, jj) * self.params.Cpgo.symbolic_expressions[jj] + + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_o(i), 0, 0), val) + + for j in range(nvar[1]): + val = 0 + for jj in range(nvar[3]): + val += U_inv[i, jj] * bips.W(jj, j) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_o(i), self._theta_a(j), 0), val * 2 * atp.sc.symbol * self.params.Lpgo.symbolic_expression) + if self.params.sbpa is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_o(i), self._theta_a(j), 0), val * self.params.sbpa.symbolic_expression) + + for j in range(nvar[3]): + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_o(i), self._deltaT_o(j), 0), - self.params.Lpgo.symbolic_expression * _kronecker_delta(i, j)) + if self.params.sbpgo is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_o(i), self._deltaT_o(j), 0), - self.params.sbpgo.symbolic_expression * _kronecker_delta(i, j)) + + for j in range(nvar[2]): + for k in range(offset, nvar[3]): + jo = j + offset # skipping the T 0 variable if it exists + + val = 0 + for jj in range(nvar[3]): + val -= U_inv[i, jj] * bips.O(jj, jo, k) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_o(i), self._psi_o(j), self._deltaT_o(k)), val) + + # deltaT_g part + if ground_temp: + for i in range(nvar[2]): + val = 0 + for jj in range(nvar[1]): + for kk in range(nvar[2]): + val += U_inv[i, kk] * bips.W(kk, jj) * self.params.Cpgo.symbolic_expressions[jj] + + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_g(i), 0, 0), val) + + for j in range(nvar[1]): + val = 0 + for jj in range(nvar[2]): + val += U_inv[i, jj] * bips.W(jj, j) + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_g(i), self._theta_a(j), 0), val * 2 * atp.sc.symbol * self.params.Lpgo.symbolic_expression) + + if self.params.sbpa is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_g(i), self._theta_a(j), 0), val * self.params.sbpa.symbolic_expression) + + for j in range(nvar[2]): + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_g(i), self._deltaT_g(j), 0), - self.params.Lpgo.symbolic_expression * _kronecker_delta(i, j)) + if self.params.sbpgo is not None: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_g(i), self._deltaT_g(j), 0), - self.params.sbpgo.symbolic_expression * _kronecker_delta(i, j)) + + return sy_arr_dic + + def compute_tensor(self): + """Routine to compute the tensor.""" + sy_arr_dic = self._compute_tensor_dicts() + sy_arr_dic = self.remove_dic_zeros(sy_arr_dic) + + if sy_arr_dic is not None: + self._set_tensor(sy_arr_dic) + + def _set_tensor(self, dic, set_symbolic=False): + self.jac_dic = self.remove_dic_zeros(self.jacobian_from_dict(dic)) + self.tensor_dic = self.remove_dic_zeros(self.simplify_dict(dic)) + + if set_symbolic: + self._set_symbolic_tensor() + + def _set_symbolic_tensor(self): + ndim = self.params.ndim + + if self.params.dynamic_T: + if self.params.T4: + raise ValueError("Symbolic tensor output not configured for T4 version, use Dynamic T version") + else: + dims = (ndim + 1, ndim + 1, ndim + 1, ndim + 1, ndim + 1) + else: + dims = (ndim + 1, ndim + 1, ndim + 1) + + jacobian_tensor = ImmutableSparseNDimArray(self.jac_dic.copy(), dims) + tensor = ImmutableSparseNDimArray(self.tensor_dic.copy(), dims) + + self.jacobian_tensor = jacobian_tensor.applyfunc(simplify) + self.tensor = tensor.applyfunc(simplify) + + @staticmethod + def remove_dic_zeros(dic): + """Removes zero values from dictionary + + Parameters + ---------- + tensor: dict + dictionary which could include 0 in values + Returns + ------- + ten_out: dict + dictionary with same keys and values as input, but keys with value of 0 are removed + """ + + non_zero_dic = dict() + for key in dic.keys(): + if dic[key] != 0: + non_zero_dic[key] = dic[key] + + return non_zero_dic + + @staticmethod + def jacobian_from_dict(dic): + """Calculates the Jacobian from the qgs tensor + + Parameters + ---------- + dic: dict + dictionary of tendencies of the model + Returns + ------- + dic_jac: dict + Jacobian tensor stored in a dictionary + """ + + rank = max([len(i) for i in dic.keys()]) + n_perm = rank - 2 + + orig_order = [i for i in range(rank)] + + keys = dic.keys() + dic_jac = dic.copy() + + for i in range(1, n_perm+1): + new_pos = orig_order.copy() + new_pos[1] = orig_order[i+1] + new_pos[i+1] = orig_order[1] + for key in keys: + dic_jac = add_to_dict(dic_jac, tuple(key[i] for i in new_pos), dic[key]) + + return dic_jac + + @staticmethod + def simplify_dict(dic): + """calculates the upper triangular tensor of a given tensor, stored in dictionary + + Parameters + ---------- + dic: dict + dictionary of tendencies of the model + + Returns + ------- + dic_upp: dict + Upper triangular tensor, stored as a tensor where the keys are the coordinates of the corresponding value. + """ + + keys = dic.keys() + dic_upp = dict() + + for key in keys: + new_key = tuple([key[0]] + sorted(key[1:])) + dic_upp = add_to_dict(dic_upp, new_key, dic[key]) + + return dic_upp + + def save_to_file(self, filename, **kwargs): + """Function to save the tensor object to a file with the :mod:`pickle` module. + + Parameters + ---------- + filename: str + The file name where to save the tensor object. + kwargs: dict + Keyword arguments to pass to the :mod:`pickle` module method. + """ + f = open(filename, 'wb') + pickle.dump(self.__dict__, f, **kwargs) + f.close() + + def sub_tensor(self, tensor=None, continuation_variables=None): + """Uses sympy substitution to convert the symbolic tensor or a symbolic dictionary to a numerical one. + + Parameters + ---------- + tensor: dict or ~sympy.tensor.array.ImmutableSparseNDimArray + + continuation_variables: Iterable(Parameter, ScalingParameter, ParametersArray) + Variables which remain symbolic, all other variables are substituted with numerical values. + If `None` all variables are substituted. + + Returns + ------- + ten_out: dict + Dictionary of the substituted tensor of the model tendencies, with coordinates and value + """ + + if continuation_variables is None: + continuation_variables = list() + + param_subs = _parameter_substitutions(self.params, continuation_variables) + + if tensor is None: + ten = self.tensor_dic + else: + ten = tensor + + if isinstance(ten, dict): + ten_out = dict() + for key in ten.keys(): + val = ten[key].subs(param_subs) + try: + ten_out[key] = float(val) + except: + ten_out[key] = val + + else: + # Assuming the tensor is a sympy tensor + ten_out = ten.subs(param_subs) + + return ten_out + + def print_tensor(self, tensor=None, dict_opp=True, tol=1e-10): + """Print the non-zero coordinates of values of the tensor of the model tendencies + + Parameters + ---------- + tensor: dict or ~sympy.tensor.array.ImmutableSparseNDimArray + Tensor of model tendencies, either as a dictionary with keys of non-zero coordinates, and values of ~sympy.core.symbol.Symbol or floats, or as a + ~sympy.tensor.array.ImmutableSparseNDimArray . + dict_opp: bool + ... + tol: float + ... + """ + if tensor is None: + if dict_opp: + temp_ten = self.tensor_dic + else: + temp_ten = self.tensor + else: + temp_ten = tensor + + if isinstance(temp_ten, dict): + val_list = [(key, temp_ten[key]) for key in temp_ten.keys()] + else: + val_list = np.ndenumerate(temp_ten) + + for ix, v in val_list: + if isinstance(v, float): + bool_test = (abs(v) > tol) + else: + bool_test = (v != 0) + + if bool_test: + try: + output_val = float(v) + except: + try: + output_val = v.simplify().evalf() + except: + output_val = v + print(str(ix) + ": " + str(output_val)) + + +class SymbolicQgsTensorDynamicT(SymbolicQgsTensor): + """qgs dynamical temperature first order (linear) symbolic tendencies tensor class. + + Parameters + ---------- + params: None or QgParams, optional + The models parameters to configure the tensor. `None` to initialize an empty tensor. Default to `None`. + atmospheric_inner_products: None or AtmosphericInnerProducts, optional + The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected. + If `None`, disable the atmospheric tendencies. Default to `None`. + oceanic_inner_products: None or OceanicInnerProducts, optional + The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected. + If `None`, disable the oceanic tendencies. Default to `None`. + ground_inner_products: None or GroundInnerProducts, optional + The inner products of the ground basis functions on which the model's PDE ground equations are projected. + If `None`, disable the ground tendencies. Default to `None`. + + Attributes + ---------- + params: None or QgParams + The models parameters used to configure the tensor. `None` for an empty tensor. + atmospheric_inner_products: None or AtmosphericInnerProducts + The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected. + If `None`, disable the atmospheric tendencies. Default to `None`. + oceanic_inner_products: None or OceanicInnerProducts + The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected. + If `None`, disable the oceanic tendencies. Default to `None`. + ground_inner_products: None or GroundInnerProducts + The inner products of the ground basis functions on which the model's PDE ground equations are projected. + If `None`, disable the ground tendencies. Default to `None`. + tensor: sparse.COO(float) + The tensor :math:`\\mathcal{T}_{i,j,k}` :math:`i`-th components. + jacobian_tensor: sparse.COO(float) + The jacobian tensor :math:`\\mathcal{T}_{i,j,k} + \\mathcal{T}_{i,k,j}` :math:`i`-th components. + """ + + def __init__(self, params=None, atmospheric_inner_products=None, oceanic_inner_products=None, ground_inner_products=None): + + SymbolicQgsTensor.__init__(self, params, atmospheric_inner_products, oceanic_inner_products, ground_inner_products) + + if params.dynamic_T: + self.compute_tensor() + + def _compute_tensor_dicts(self): + + if self.params is None: + return None + + if self.atmospheric_inner_products is None and self.oceanic_inner_products is None \ + and self.ground_inner_products is None: + return None + + aips = self.atmospheric_inner_products + + bips = None + if self.oceanic_inner_products is not None: + bips = self.oceanic_inner_products + + elif self.ground_inner_products is not None: + bips = self.ground_inner_products + + if bips is not None: + go = bips.stored + else: + go = True + + if aips.stored and go: + symbolic_array_full_dict = self._compute_stored_full_dict() + + else: + symbolic_array_full_dict = self._compute_non_stored_full_dict() + + return symbolic_array_full_dict + + def _compute_stored_full_dict(self): + par = self.params + nvar = par.number_of_variables + aips = self.atmospheric_inner_products + + bips = None + if self.oceanic_inner_products is not None: + bips = self.oceanic_inner_products + ocean = True + else: + ocean = False + + if self.ground_inner_products is not None: + bips = self.ground_inner_products + ground_temp = True + else: + ground_temp = False + + # constructing some derived matrices + if aips is not None: + + a_theta = dict() + for i in range(nvar[1]): + for j in range(nvar[1]): + a_theta[(i, j)] = par.atmospheric_params.sig0.symbol * aips.a(i, j) - aips.u(i, j) + a_theta = ImmutableSparseMatrix(nvar[1], nvar[1], a_theta) + a_theta = a_theta.inverse() + + if bips is not None: + U_inv = dict() + if ocean: + for i in range(nvar[3]): + for j in range(nvar[3]): + U_inv[(i, j)] = bips.U(i, j) + U_inv = ImmutableSparseMatrix(nvar[3], nvar[3], U_inv) + U_inv = U_inv.inverse() + else: + for i in range(nvar[2]): + for j in range(nvar[2]): + U_inv[i, j] = bips.U(i, j) + U_inv = ImmutableSparseMatrix(nvar[2], nvar[2], U_inv) + U_inv = U_inv.inverse() + + ################# + + sy_arr_dic = dict() + # theta_a part + for i in range(nvar[1]): + if self.params.T4LSBpa is not None: + j = k = ell = 0 + for m in range(nvar[1]): + val = 0 + for jj in range(nvar[1]): + val += a_theta[i, jj] * aips._z[jj, j, k, ell, m] + if m == 0: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._theta_a(j), self._theta_a(k), self._theta_a(ell), self._theta_a(m)), self.params.T4LSBpa.symbolic_expression * val) + else: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._theta_a(j), self._theta_a(k), self._theta_a(ell), self._theta_a(m)), 4 * self.params.T4LSBpa.symbolic_expression * val) + + if ocean: + if self.params.T4LSBpgo is not None: + j = k = ell = 0 + for m in range(nvar[3]): + val = 0 + for jj in range(nvar[1]): + val += a_theta[i, jj] * aips._v[jj, j, k, ell, m] + if m == 0: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._deltaT_o(j), self._deltaT_o(k), self._deltaT_o(ell), self._deltaT_o(m)), -self.params.T4LSBpgo.symbolic_expression * val) + else: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._deltaT_o(j), self._deltaT_o(k), self._deltaT_o(ell), self._deltaT_o(m)), -4 * self.params.T4LSBpgo.symbolic_expression * val) + + if ground_temp: + if self.params.T4LSBpgo is not None: + j = k = ell = 0 + for m in range(nvar[2]): + val = 0 + for jj in range(nvar[1]): + val += a_theta[i, jj] * aips._v[jj, j, k, ell, m] + if m == 0: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._deltaT_g(j), self._deltaT_g(k), self._deltaT_g(ell), self._deltaT_g(m)), -self.params.T4LSBpgo.symbolic_expression * val) + else: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._theta_a(i), self._deltaT_g(j), self._deltaT_g(k), self._deltaT_g(ell), self._deltaT_g(m)), -4 * self.params.T4LSBpgo.symbolic_expression * val) + + if ocean: + # delta_T part + for i in range(nvar[3]): + j = k = ell = 0 + for m in range(nvar[1]): + val = 0 + for jj in range(nvar[3]): + val += U_inv[i, jj] * bips._Z[jj, j, k, ell, m] + if m == 0: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_o(i), self._theta_a(j), self._theta_a(k), self._theta_a(ell), self._theta_a(m)), self.params.T4sbpa.symbolic_expression * val) + else: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_o(i), self._theta_a(j), self._theta_a(k), self._theta_a(ell), self._theta_a(m)), 4 * self.params.T4sbpa.symbolic_expression * val) + + for m in range(nvar[3]): + val = 0 + for jj in range(nvar[3]): + val += U_inv[i, jj] * bips._V[jj, j, k, ell, m] + if m == 0: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_o(i), self._deltaT_o(j), self._deltaT_o(k), self._deltaT_o(ell), self._deltaT_o(m)), - self.params.T4sbpgo.symbolic_expression * val) + else: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_o(i), self._deltaT_o(j), self._deltaT_o(k), self._deltaT_o(ell), self._deltaT_o(m)), -4 * self.params.T4sbpgo.symbolic_expression * val) + + if ground_temp: + # deltaT_g part + for i in range(nvar[2]): + j = k = ell = 0 + for m in range(nvar[1]): + val = 0 + for jj in range(nvar[2]): + val += U_inv[i, jj] * bips._Z[jj, j, k, ell, m] + if m == 0: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_g(i), self._theta_a(j), self._theta_a(k), self._theta_a(ell), self._theta_a(m)), self.params.T4sbpa.symbolic_expression * val) + else: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_g(i), self._theta_a(j), self._theta_a(k), self._theta_a(ell), self._theta_a(m)), 4 * self.params.T4sbpa.symbolic_expression * val) + + for m in range(nvar[2]): + val = 0 + for jj in range(nvar[2]): + val += U_inv[i, jj] * bips._V[jj, j, k, ell, m] + if m == 0: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_g(i), self._deltaT_g(j), self._deltaT_g(k), self._deltaT_g(ell), self._deltaT_g(m)), -self.params.T4sbpgo.symbolic_expression * val) + else: + sy_arr_dic = add_to_dict(sy_arr_dic, (self._deltaT_g(i), self._deltaT_g(j), self._deltaT_g(k), self._deltaT_g(ell), self._deltaT_g(m)), -4 * self.params.T4sbpgo.symbolic_expression * val) + + return sy_arr_dic + + def _compute_non_stored_full_dict(self): + par = self.params + nvar = par.number_of_variables + aips = self.atmospheric_inner_products + + bips = None + if self.oceanic_inner_products is not None: + bips = self.oceanic_inner_products + ocean = True + else: + ocean = False + + if self.ground_inner_products is not None: + bips = self.ground_inner_products + ground_temp = True + else: + ground_temp = False + + # constructing some derived matrices + if aips is not None: + a_theta = dict() + for i in range(nvar[1]): + for j in range(nvar[1]): + a_theta[(i, j)] = par.atmospheric_params.sig0.symbol * aips.a(i, j) - aips.u(i, j) + + a_theta = ImmutableSparseMatrix(nvar[1], nvar[1], a_theta) + a_theta = a_theta.inverse() + + if bips is not None: + if ocean: + U_inv = dict() + for i in range(nvar[3]): + for j in range(nvar[3]): + U_inv[i, j] = bips.U(i, j) + U_inv = ImmutableSparseMatrix(nvar[3], nvar[3], U_inv) + + else: + U_inv = dict() + for i in range(nvar[2]): + for j in range(nvar[2]): + U_inv[(i, j)] = bips.U(i, j) + U_inv = ImmutableSparseMatrix(nvar[2], nvar[2], U_inv) + + U_inv = U_inv.inverse() + + ################# + + sy_arr_dic = dict() + # theta_a part + for i in range(nvar[1]): + + if self.params.T4LSBpa is not None: + j = k = ell = 0 + for m in range(nvar[1]): + val = 0 + for jj in range(nvar[1]): + val += a_theta[i, jj] * aips.z(jj, j, k, ell, m) + if m == 0: + sy_arr_dic[(self._theta_a(i), self._theta_a(j), self._theta_a(k), self._theta_a(ell), self._theta_a(m))] = self.params.T4LSBpa.symbolic_expression * val + else: + sy_arr_dic[(self._theta_a(i), self._theta_a(j), self._theta_a(k), self._theta_a(ell), self._theta_a(m))] = 4 * self.params.T4LSBpa.symbolic_expression * val + + if ocean: + if self.params.T4LSBpgo is not None: + j = k = ell = 0 + for m in range(nvar[3]): + val = 0 + for jj in range(nvar[1]): + val -= a_theta[i, jj] * aips.v(jj, j, k, ell, m) + if m == 0: + sy_arr_dic[(self._theta_a(i), self._deltaT_o(j), self._deltaT_o(k), self._deltaT_o(ell), self._deltaT_o(m))] = self.params.T4LSBpgo.symbolic_expression * val + else: + sy_arr_dic[(self._theta_a(i), self._deltaT_o(j), self._deltaT_o(k), self._deltaT_o(ell), self._deltaT_o(m))] = 4 * self.params.T4LSBpgo.symbolic_expression * val + + if ground_temp: + if self.params.T4LSBpgo is not None: + j = k = ell = 0 + for m in range(nvar[2]): + val = 0 + for jj in range(nvar[1]): + val -= a_theta[i, jj] * aips.v(jj, j, k, ell, m) + if m == 0: + sy_arr_dic[(self._theta_a(i), self._deltaT_g(j), self._deltaT_g(k), self._deltaT_g(ell), self._deltaT_g(m))] = self.params.T4LSBpgo.symbolic_expression * val + else: + sy_arr_dic[(self._theta_a(i), self._deltaT_g(j), self._deltaT_g(k), self._deltaT_g(ell), self._deltaT_g(m))] = 4 * self.params.T4LSBpgo.symbolic_expression * val + + if ocean: + + # deltaT_o part + for i in range(nvar[3]): + + j = k = ell = 0 + for m in range(nvar[1]): + val = 0 + for jj in range(nvar[3]): + val += U_inv[i, jj] * bips.Z(jj, j, k, ell, m) + if m == 0: + sy_arr_dic[(self._deltaT_o(i), self._theta_a(j), self._theta_a(k), self._theta_a(ell), self._theta_a(m))] = self.params.T4sbpa.symbolic_expression * val + else: + sy_arr_dic[(self._deltaT_o(i), self._theta_a(j), self._theta_a(k), self._theta_a(ell), self._theta_a(m))] = 4 * self.params.T4sbpa.symbolic_expression * val + + for m in range(nvar[3]): + val = 0 + for jj in range(nvar[3]): + val -= U_inv[i, jj] * bips.V(jj, j, k, ell, m) + if m == 0: + sy_arr_dic[(self._deltaT_o(i), self._deltaT_o(j), self._deltaT_o(k), self._deltaT_o(ell), self._deltaT_o(m))] = self.params.T4sbpgo.symbolic_expression * val + else: + sy_arr_dic[(self._deltaT_o(i), self._deltaT_o(j), self._deltaT_o(k), self._deltaT_o(ell), self._deltaT_o(m))] = 4 * self.params.T4sbpgo.symbolic_expression * val + + # deltaT_g part + if ground_temp: + for i in range(nvar[2]): + + j = k = ell = 0 + for m in range(nvar[1]): + val = 0 + for jj in range(nvar[2]): + val += U_inv[i, jj] * bips._Z[jj, j, k, ell, m] + if m == 0: + sy_arr_dic[(self._deltaT_g(i), self._theta_a(j), self._theta_a(k), self._theta_a(ell), self._theta_a(m))] = self.params.T4sbpa.symbolic_expression * val + else: + sy_arr_dic[(self._deltaT_g(i), self._theta_a(j), self._theta_a(k), self._theta_a(ell), self._theta_a(m))] = 4 * self.params.T4sbpa.symbolic_expression * val + + for m in range(nvar[2]): + val = 0 + for jj in range(nvar[2]): + val -= U_inv[i, jj] * bips._V[jj, j, k, ell, m] + if m == 0: + sy_arr_dic[(self._deltaT_g(i), self._deltaT_g(j), self._deltaT_g(k), self._deltaT_g(ell), self._deltaT_g(m))] = self.params.T4sbpgo.symbolic_expression * val + else: + sy_arr_dic[(self._deltaT_g(i), self._deltaT_g(j), self._deltaT_g(k), self._deltaT_g(ell), self._deltaT_g(m))] = 4 * self.params.T4sbpgo.symbolic_expression * val + + return sy_arr_dic + + def compute_tensor(self): + """Routine to compute the tensor.""" + # gathering + if self.params.T4: + # TODO: Make a proper error message for here + raise ValueError("Parameters are set for T4 version, set dynamic_T=True") + + symbolic_dict_linear = SymbolicQgsTensor._compute_tensor_dicts(self) + symbolic_dict_linear = _shift_dict_keys(symbolic_dict_linear, (0, 0)) + + symbolic_dict_dynT = self._compute_tensor_dicts() + + if symbolic_dict_linear is not None: + symbolic_dict_dynT = {**symbolic_dict_linear, **symbolic_dict_dynT} + + if symbolic_dict_dynT is not None: + self._set_tensor(symbolic_dict_dynT) + + +class SymbolicQgsTensorT4(SymbolicQgsTensor): + # TODO: this takes a long time (>1hr) to run. I think we need a better way to run the non-stored z, v, Z, V IPs. Maybe do not allow `n` as a continuation parameter for this version? + # TODO: Create a warning about long run-times. + + """qgs dynamical temperature first order (linear) symbolic tendencies tensor class. + + Parameters + ---------- + params: None or QgParams, optional + The models parameters to configure the tensor. `None` to initialize an empty tensor. Default to `None`. + atmospheric_inner_products: None or AtmosphericInnerProducts, optional + The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected. + If `None`, disable the atmospheric tendencies. Default to `None`. + oceanic_inner_products: None or OceanicInnerProducts, optional + The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected. + If `None`, disable the oceanic tendencies. Default to `None`. + ground_inner_products: None or GroundInnerProducts, optional + The inner products of the ground basis functions on which the model's PDE ground equations are projected. + If `None`, disable the ground tendencies. Default to `None`. + + Attributes + ---------- + params: None or QgParams + The models parameters used to configure the tensor. `None` for an empty tensor. + atmospheric_inner_products: None or AtmosphericInnerProducts + The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected. + If `None`, disable the atmospheric tendencies. Default to `None`. + oceanic_inner_products: None or OceanicInnerProducts + The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected. + If `None`, disable the oceanic tendencies. Default to `None`. + ground_inner_products: None or GroundInnerProducts + The inner products of the ground basis functions on which the model's PDE ground equations are projected. + If `None`, disable the ground tendencies. Default to `None`. + tensor: sparse.COO(float) + The tensor :math:`\\mathcal{T}_{i,j,k}` :math:`i`-th components. + jacobian_tensor: sparse.COO(float) + The jacobian tensor :math:`\\mathcal{T}_{i,j,k} + \\mathcal{T}_{i,k,j}` :math:`i`-th components. + """ + + def __init__(self, params=None, atmospheric_inner_products=None, oceanic_inner_products=None, ground_inner_products=None): + + SymbolicQgsTensor.__init__(self, params, atmospheric_inner_products, oceanic_inner_products, ground_inner_products) + + if params.T4: + self.compute_tensor() + + def _compute_non_stored_full_dict(self): + par = self.params + nvar = par.number_of_variables + aips = self.atmospheric_inner_products + + bips = None + if self.oceanic_inner_products is not None: + bips = self.oceanic_inner_products + ocean = True + else: + ocean = False + + if self.ground_inner_products is not None: + bips = self.ground_inner_products + ground_temp = True + else: + ground_temp = False + + # constructing some derived matrices + if aips is not None: + a_theta = dict() + for i in range(nvar[1]): + for j in range(nvar[1]): + a_theta[(i, j)] = par.atmospheric_params.sig0.symbol * aips.a(i, j) - aips.u(i, j) + + a_theta = ImmutableSparseMatrix(nvar[1], nvar[1], a_theta) + a_theta = a_theta.inverse() + + if bips is not None: + if ocean: + U_inv = dict() + for i in range(nvar[3]): + for j in range(nvar[3]): + U_inv[i, j] = bips.U(i, j) + U_inv = ImmutableSparseMatrix(nvar[3], nvar[3], U_inv) + + else: + U_inv = dict() + for i in range(nvar[2]): + for j in range(nvar[2]): + U_inv[(i, j)] = bips.U(i, j) + U_inv = ImmutableSparseMatrix(nvar[2], nvar[2], U_inv) + + U_inv = U_inv.inverse() + + ################# + + sy_arr_dic = dict() + # theta_a part + for i in range(nvar[1]): + + if self.params.T4LSBpa is not None: + for j in range(nvar[1]): + for k in range(nvar[1]): + for ell in range(nvar[1]): + for m in range(nvar[1]): + val = 0 + for jj in range(nvar[1]): + val += a_theta[i, jj] * aips.z(jj, j, k, ell, m) + + sy_arr_dic[(self._theta_a(i), self._theta_a(j), self._theta_a(k), self._theta_a(ell), self._theta_a(m))] = self.params.T4LSBpa.symbolic_expression * val + + if ocean: + if self.params.T4LSBpgo is not None: + for j in range(nvar[3]): + for k in range(nvar[3]): + for ell in range(nvar[3]): + for m in range(nvar[3]): + val = 0 + for jj in range(nvar[1]): + val -= a_theta[i, jj] * aips.v(jj, j, k, ell, m) + + sy_arr_dic[(self._theta_a(i), self._deltaT_o(j), self._deltaT_o(k), self._deltaT_o(ell), self._deltaT_o(m))] = self.params.T4LSBpgo.symbolic_expression * val + + if ground_temp: + if self.params.T4LSBpgo is not None: + for j in range(nvar[2]): + for k in range(nvar[2]): + for ell in range(nvar[2]): + for m in range(nvar[2]): + val = 0 + for jj in range(nvar[1]): + val -= a_theta[i, jj] * aips.v(jj, j, k, ell, m) + + sy_arr_dic[(self._theta_a(i), self._deltaT_g(j), self._deltaT_g(k), self._deltaT_g(ell), self._deltaT_g(m))] = self.params.T4LSBpgo.symbolic_expression * val + + if ocean: + + # deltaT_o part + for i in range(nvar[3]): + for j in range(nvar[1]): + for k in range(nvar[1]): + for ell in range(nvar[1]): + for m in range(nvar[1]): + val = 0 + for jj in range(nvar[3]): + val += U_inv[i, jj] * bips.Z(jj, j, k, ell, m) + + sy_arr_dic[(self._deltaT_o(i), self._theta_a(j), self._theta_a(k), self._theta_a(ell), self._theta_a(m))] = self.params.T4LSBpa.symbolic_expression * val + + for j in range(nvar[3]): + for k in range(nvar[3]): + for ell in range(nvar[3]): + for m in range(nvar[3]): + val = 0 + for jj in range(nvar[3]): + val -= U_inv[i, jj] * bips.V(jj, j, k, ell, m) + + sy_arr_dic[(self._deltaT_o(i), self._deltaT_o(j), self._deltaT_o(k), self._deltaT_o(ell), self._deltaT_o(m))] = self.params.T4LSBpgo.symbolic_expression * val + + # deltaT_g part + if ground_temp: + for i in range(nvar[2]): + for j in range(nvar[1]): + for k in range(nvar[1]): + for ell in range(nvar[1]): + for m in range(nvar[1]): + val = 0 + for jj in range(nvar[2]): + val += U_inv[i, jj] * bips._Z[jj, j, k, ell, m] + + sy_arr_dic[(self._deltaT_g(i), self._theta_a(j), self._theta_a(k), self._theta_a(ell), self._theta_a(m))] = self.params.T4LSBpa.symbolic_expression * val + + for j in range(nvar[2]): + for k in range(nvar[2]): + for ell in range(nvar[2]): + for m in range(nvar[2]): + val = 0 + for jj in range(nvar[2]): + val -= U_inv[i, jj] * bips._V[jj, j, k, ell, m] + + sy_arr_dic[(self._deltaT_g(i), self._deltaT_g(j), self._deltaT_g(k), self._deltaT_g(ell), self._deltaT_g(m))] = self.params.T4LSBpgo.symbolic_expression * val + + return sy_arr_dic + + def compute_tensor(self): + """Routine to compute the tensor.""" + # gathering + if not self.params.T4: + raise ValueError("Parameters are not set for T4 version") + + symbolic_dict_linear = SymbolicQgsTensor._compute_tensor_dicts(self) + symbolic_dict_linear = _shift_dict_keys(symbolic_dict_linear, (0, 0)) + + symbolic_dict_T4 = self._compute_non_stored_full_dict() + + if symbolic_dict_linear is not None: + symbolic_dict_T4 = {**symbolic_dict_linear, **symbolic_dict_T4} + + if symbolic_dict_T4 is not None: + self._set_tensor(symbolic_dict_T4) + + +def _kronecker_delta(i, j): + + if i == j: + return 1 + + else: + return 0 + + +def _shift_dict_keys(dic, shift): + """ + Keys of given dictionary are altered to add values in the given indicies + + Parameters + ---------- + dic: dictionary + + shift: Tuple + """ + + shifted_dic = dict() + for key in dic.keys(): + new_key = key + shift + shifted_dic[new_key] = dic[key] + + return shifted_dic + + +def _parameter_substitutions(params, continuation_variables): + + subs = _parameter_values(params) + for _, obj in params.__dict__.items(): + if issubclass(obj.__class__, Params): + subs.update(_parameter_values(obj)) + + # Manually add properties from class + subs[params.scale_params.L.symbol] = params.scale_params.L + subs[params.scale_params.beta.symbol] = params.scale_params.beta + + # Remove variables in continuation variables + for cv in continuation_variables: + if isinstance(cv, ParametersArray): + for cv_i in cv.symbols: + subs.pop(cv_i) + elif hasattr(cv, "symbol"): + subs.pop(cv.symbol) + else: # Try ... who knows... + subs.pop(cv) + + return subs + + +def _parameter_values(pars): + """Function takes a parameter class and produces a dictionary of the symbol and the corresponding numerical value""" + + subs = dict() + for val in pars.__dict__.values(): + if isinstance(val, Parameter): + if val.symbol is not None: + subs[val.symbol] = val + + if isinstance(val, ScalingParameter): + if val.symbol is not None: + subs[val.symbol] = val + + if isinstance(val, ParametersArray): + for v in val: + if v.symbol is not None or v.symbol != 0: + subs[v.symbol] = v + return subs + + +if __name__ == "__main__": + dic = dict() + dic = add_to_dict(dic, (0, 0), 1) + dic = add_to_dict(dic, (0, 0), 2) + print(dic) + + from qgs.params.params import QgParams + from qgs.inner_products import symbolic + + params = QgParams({'rr': 287.e0, 'sb': 5.6e-8}) + params.set_atmospheric_channel_fourier_modes(6, 6, mode="symbolic") + params.atmospheric_params.set_params({'sigma': 0.2, 'kd': 0.1, 'kdp': 0.01}) + + params.ground_params.set_orography(0.2, 1) + params.atemperature_params.set_thetas(0.1, 0) + + aip = symbolic.AtmosphericSymbolicInnerProducts(params, return_symbolic=True, make_substitution=True) + + # sym_aotensor = SymbolicQgsTensor(params=params, atmospheric_inner_products=aip)