Merge branch 'main' into qgs

README.md

@@ -8,6 +8,15 @@ The field of data assimilation (DA) studies the integration of theory with obser
 
 Today, applied DA has matured in operational weather forecasting to include the entire online cycled process of continually ingesting numerous disparate observational data sources and integrating them with numerical prediction models to make regular forecasts, while estimating errors and uncertainties in this process and accounting for them along the way. The process can also include correcting inaccuracies in the model formulations or applying post-processing to forecasts to improve agreement with observations.  
 
+Most of the software here was adapted from code developed for the following publications. Please cite these works when using this software package: 
+
+- Penny, S. G., Smith, T. A., Chen, T.-C., Platt, J. A., Lin, H.-Y., Goodliff, M., & Abarbanel, H. D. I. (2022). Integrating recurrent neural networks with data assimilation for scalable data-driven state estimation. Journal of Advances in Modeling Earth Systems, 14, e2021MS002843. https://doi.org/10.1029/2021MS002843  
+
+- Smith, T. A., Penny, S. G., Platt, J. A., & Chen, T.-C. (2023). Temporal subsampling diminishes small spatial scales in recurrent neural network emulators of geophysical turbulence. Journal of Advances in Modeling Earth Systems, 15, e2023MS003792. https://doi.org/10.1029/2023MS003792  
+
+- Platt, J.A., S.G. Penny, T.A. Smith, T.-C. Chen, H.D.I. Abarbanel, (2022). A systematic exploration of reservoir computing for forecasting complex spatiotemporal dynamics,
+Neural Networks,Volume 153, 530-552, ISSN 0893-6080, https://doi.org/10.1016/j.neunet.2022.06.025.
+
 
 ## Installation
 

dabench/data/_utils.py

@@ -30,23 +30,19 @@ def integrate(function, x0, t_final, delta_t, method='odeint', stride=None,
         (time_dim, system_dim) and t is time array with shape (time_dim)
     """
     if method == 'odeint':
+        # Define timesteps
         if jax_comps:
-            # Define timesteps
             t = jnp.arange(0.0, t_final, delta_t)
-            # If stride is defined, remove steps that are not on stride steps
-            if stride is not None:
-                assert stride > 1 and isinstance(stride, int), \
-                    'integrate: stride = {}, must be > 1 and an int'.format(
-                            stride)
-            t = t[::stride]
-            y = odeint(function, x0, t, **kwargs)
         else:
             t = np.arange(0.0, t_final, delta_t)
-            if stride is not None:
-                assert stride > 1 and isinstance(stride, int), \
-                    'integrate: stride = {}, must be > 1 and an int'.format(
-                            stride)
+        # If stride is defined, remove timesteps that are not on stride steps
+        if stride is not None:
+            assert stride > 1 and isinstance(stride, int), \
+                'integrate: stride = {}, must be > 1 and an int'.format(stride)
             t = t[::stride]
+        if jax_comps:
+            y = odeint(function, x0, t, **kwargs)
+        else:
             y = spodeint(function, x0, t, **kwargs)
     else:
         raise 'integration method {} not supported'.format(method)

tests/data_aws_test.py

@@ -5,7 +5,7 @@
 import numpy as np
 
 
-@pytest.fixture
+@pytest.fixture(scope='module')
 def aws_small():
     """Defines aws object for rest of tests"""
     aws_obj = AWS(months=['01', '02'])
@@ -14,7 +14,7 @@ def aws_small():
     return aws_obj
 
 
-@pytest.fixture
+@pytest.fixture(scope='module')
 def aws_multivar():
     """Defines aws object for rest of tests"""
     aws_obj = AWS(variables=['air_temperature_at_2_metres',

tests/data_barotropic_test.py

@@ -5,7 +5,7 @@
 import pytest
 
 
-@pytest.fixture
+@pytest.fixture(scope='module')
 def barotropic():
     """Defines class Barotropic object for rest of tests."""
     barotropic_obj = Barotropic()

tests/data_gcp_test.py

@@ -8,7 +8,7 @@
 import numpy as np
 
 
-@pytest.fixture
+@pytest.fixture(scope='module')
 def gcp_small():
     """Defines gcp object for rest of tests"""
     gcp_obj = GCP(data_type='single-level-reanalysis',
@@ -18,7 +18,7 @@ def gcp_small():
     return gcp_obj
 
 
-@pytest.fixture
+@pytest.fixture(scope='module')
 def gcp_multivar():
     """Defines gcp object for rest of tests"""
     gcp_obj = GCP(data_type='single-level-forecast',

tests/data_pyqg_test.py

@@ -6,7 +6,7 @@
 from dabench.data import PyQG
 
 
-@pytest.fixture
+@pytest.fixture(scope='module')
 def pyqg():
     """Defines class PYQ object for rest of tests."""
     pyqg_obj = PyQG()