moved some methods to PWCModelRunFD from ModelDataObject in preparati…

…on of saving models to netcdf

CompartmentalSystems/pwc_model_run_fd.py

@@ -10,7 +10,7 @@
 from .pwc_model_run import PWCModelRun
 from .smooth_reservoir_model import SmoothReservoirModel
 from .myOdeResult import solve_ivp_pwc
-
+from bgc_md2.Variable import Variable
 
 ###############################################################################
 
@@ -140,6 +140,85 @@ def acc_net_external_output_vector(self):
     def fake_discretized_Bs(self, data_times=None):
         return self.pwc_mr.fake_discretized_Bs(data_times)
 
+    def get_stock(self, mdo, pool_name, nr_layer=0, name=''):
+        ms = mdo.model_structure
+        pool_nr = ms.get_pool_nr(pool_name, nr_layer)
+        soln = self.solve()
+
+        if name == '':
+            name = ms.get_stock_var(pool_name)
+
+        return Variable(
+            name=name,
+            data=soln[:, pool_nr],
+            unit=mdo.stock_unit
+        )
+
+    def get_acc_gross_external_input_flux(self, mdo, pool_name, nr_layer=0, name=''):
+        ms = mdo.model_structure
+        pool_nr = ms.get_pool_nr(pool_name, nr_layer)
+
+        if name == '':
+            name = ms.get_external_input_flux_var(pool_name)
+
+        return Variable(
+            name=name,
+            data=self.acc_gross_external_input_vector()[:, pool_nr],
+            unit=mdo.stock_unit
+        )
+
+    def get_acc_gross_external_output_flux(self, mdo, pool_name, nr_layer=0, name=''):
+        ms = mdo.model_structure
+        pool_nr = ms.get_pool_nr(pool_name, nr_layer)
+
+        if name == '':
+            name = ms.get_external_output_flux_var(pool_name)
+
+        return Variable(
+            name=name,
+            data=self.acc_gross_external_output_vector()[:, pool_nr],
+            unit=mdo.stock_unit
+        )
+
+    def get_acc_gross_internal_flux(
+        self,
+        mdo,
+        pool_name_from,
+        pool_name_to,
+        nr_layer_from=0,
+        nr_layer_to=0,
+        name=''
+    ):
+        ms = mdo.model_structure
+        pool_nr_from = ms.get_pool_nr(
+            pool_name_from,
+            nr_layer_from
+        )
+        pool_nr_to = ms.get_pool_nr(
+            pool_name_to,
+            nr_layer_to
+        )
+        Fs = self.acc_gross_internal_flux_matrix()
+
+        if name == '':
+            name = ms.get_horizontal_flux_var(pool_name_from, pool_name_to)
+
+        return Variable(
+            name=name,
+            data=Fs[:, pool_nr_to, pool_nr_from],
+            unit=mdo.stock_unit
+        )
+
+#    def to_netcdf(self, mdo, file_path):
+#        data_vars = {}
+#        model_ds = xr.Dataset(
+#            data_vars=data_vars,
+#            coords=coords,
+#            attr=attrs
+#        }
+#        model_ds.to_netcdf(file_path)
+#        model_ds.close()
+
 #    @classmethod
 #    def load_from_file(cls, filename):
 #        pwc_mr_fd_dict = picklegzip.load(filename)

tests/Test_pwc_model_run_fd.py

@@ -1,13 +1,17 @@
 import unittest
 
 import numpy as np
-
+import xarray as xr
 from sympy import Function, Matrix, symbols
 
 from CompartmentalSystems.smooth_reservoir_model import SmoothReservoirModel
 from CompartmentalSystems.smooth_model_run import SmoothModelRun
 from CompartmentalSystems.pwc_model_run_fd import PWCModelRunFD
 
+from bgc_md2.models.CARDAMOM.CARDAMOMlib import load_mdo
+
+import os.path
+THIS_DIR = os.path.dirname(os.path.abspath(__file__))
 
 class TestPWCModelRunFD(unittest.TestCase):
 
@@ -88,6 +92,101 @@ def test_reconstruction_accuracy(self):
             )
         )
 
+    def test_get_stock(self):
+        filename = 'cardamom_for_holger_10_ensembles.nc'
+        my_data_path = os.path.join(THIS_DIR, filename)
+        dataset = xr.open_dataset(my_data_path)
+        ds = dataset.isel(ens=0, lat=0, lon=0, time=slice(None, 24))
+        mdo = load_mdo(ds)
+        mr = mdo.create_model_run()
+
+        soln = mr.solve()
+        for nr, pool_dict in enumerate(mdo.model_structure.pool_structure):
+            with self.subTest():
+                pool_name = pool_dict['pool_name']
+                self.assertTrue(
+                    np.all(soln[:, nr] == mr.get_stock(mdo, pool_name).data)
+                ) 
+
+    def test_get_acc_gross_external_input_flux(self):
+        filename = 'cardamom_for_holger_10_ensembles.nc'
+        my_data_path = os.path.join(THIS_DIR, filename)
+        dataset = xr.open_dataset(my_data_path)
+        ds = dataset.isel(ens=0, lat=0, lon=0, time=slice(None, 24))
+        mdo = load_mdo(ds)
+        mr = mdo.create_model_run()
+
+        ms = mdo.model_structure
+        Us = mr.acc_gross_external_input_vector()
+        for nr, pool_dict in enumerate(ms.pool_structure):
+            with self.subTest():
+                pool_name = pool_dict['pool_name']
+                if pool_name in ms.external_input_structure.keys():
+                    self.assertTrue(
+                        np.all(Us[:, nr] ==\
+                               mr.get_acc_gross_external_input_flux(
+                                   mdo, 
+                                   pool_name
+                               ).data
+                        )
+                    )
+                else:
+                    self.assertTrue(np.all(Us[:, nr] == 0))
+
+    def test_get_acc_gross_external_output_flux(self):
+        filename = 'cardamom_for_holger_10_ensembles.nc'
+        my_data_path = os.path.join(THIS_DIR, filename)
+        dataset = xr.open_dataset(my_data_path)
+        ds = dataset.isel(ens=0, lat=0, lon=0, time=slice(None, 24))
+        mdo = load_mdo(ds)
+        mr = mdo.create_model_run()
+
+        ms = mdo.model_structure
+        Rs = mr.acc_gross_external_output_vector()
+        for nr, pool_dict in enumerate(ms.pool_structure):
+            with self.subTest():
+                pool_name = pool_dict['pool_name']
+                if pool_name in ms.external_output_structure.keys():
+                    self.assertTrue(
+                        np.all(Rs[:, nr] ==\
+                               mr.get_acc_gross_external_output_flux(
+                                   mdo, 
+                                   pool_name
+                               ).data
+                        )
+                    )
+                else:
+                    self.assertTrue(np.all(Rs[:, nr] == 0))
+
+    def test_get_acc_gross_internal_flux(self):
+        filename = 'cardamom_for_holger_10_ensembles.nc'
+        my_data_path = os.path.join(THIS_DIR, filename)
+        dataset = xr.open_dataset(my_data_path)
+        ds = dataset.isel(ens=0, lat=0, lon=0, time=slice(None, 24))
+        mdo = load_mdo(ds)
+        mr = mdo.create_model_run()
+
+        ms = mdo.model_structure
+        Fs = mr.acc_gross_internal_flux_matrix()
+        for nr_from, pool_dict_from in enumerate(ms.pool_structure):
+            pool_name_from = pool_dict_from['pool_name']
+            for nr_to, pool_dict_to in enumerate(ms.pool_structure):
+                pool_name_to = pool_dict_to['pool_name']
+                with self.subTest():
+                    if (pool_name_from, pool_name_to) in\
+                        ms.horizontal_structure.keys():
+                        self.assertTrue(
+                            np.all(Fs[:, nr_to, nr_from] ==\
+                                   mr.get_acc_gross_internal_flux(
+                                       mdo, 
+                                       pool_name_from,
+                                       pool_name_to
+                                   ).data
+                            )
+                        )
+                    else:
+                        self.assertTrue(np.all(Fs[:, nr_to, nr_from] == 0))
+
 ###############################################################################