Pip packaging (#44)

* Add dependencies to pyproject.toml

* Remove environment.yml

* Replace from jax.config import config (deprecated) with just jax.config

* Streamline resources and xarray manipulation to match new versions

* Call asnumpy on xarray ds when opening to avoid scipyarraywrapper attribute error

* Replace jax.numpy in1d (no longer exists) with isin

* Update getargspec (deprecated) to inspect.signature

* Add qgs to pyproject.toml

* Fix mistake, signature() must called on h

* Add zarr, dask, and others needed for google cloud downloads

* Need to precompute indexer when using ds.where on remote zarr array

* Add qgs and cloud to dev install (will add pyqg later)

* Move from conda to pip for CI testing

* Rename CI config yml from conda to pip

* Add flake8 to dev build

* Use workaround for installing pyqg with older Cython

* Manually install wheel before pyqg install

* Add dask for cloud build

* Remove pyqg from dependencies for now, include tests in full install

* ETKF fix to exclude end of timewindow exact match observations

* Add extras to example build (ray, pandas, matplotlib)

* Add hyperopt to examples build

* Add cftime to examples (needed for loading netcdf)

* Update README with new install instructions

* Updated README with new install instructions

---------

Co-authored-by: Kylen Solvik <[email protected]>

Author: StevePny

.github/workflows/python-ci-conda.yml renamed to .github/workflows/python-ci-pip.yml

@@ -13,32 +13,19 @@ jobs:
     steps:
 
     - uses: actions/checkout@v3
-    - name: Set up Python 3.10.8
+    - name: Set up Python 3.11.9
       uses: actions/setup-python@v3
       with:
-        python-version: '3.10.8'
+        python-version: '3.11.9'
 
-    - name: Add conda to system path
-      run: |
-        # $CONDA is an environment variable pointing to the root of the miniconda directory
-        echo $CONDA/bin >> $GITHUB_PATH
     - name: Set cache date
       run: echo "DATE=$(date +'%Y%m%d')" >> $GITHUB_ENV
 
-    - uses: actions/cache@v2
-      with:
-        path: /usr/share/miniconda3/envs/base/
-        key: conda-${{ hashFiles('environment.yml') }}-${{ env.DATE }}-${{ env.CACHE_NUMBER }}
-      id: cache
-      
-    - name: Update environment
+    - name: Install package and requirements
       run: |
-        conda install conda-libmamba-solver flake8 pytest
-        # Slightly awkward way of adding an optional package to our environment
-        echo "  - pyqg" >> environment.yml
-        conda env update --name base --file environment.yml --solver=libmamba
-        pip install qgs
-      if: steps.cache.outputs.cache-hit != 'true'
+        python3 -m pip install -e ".[dev]"
+        python3 -m pip install Cython==0.29.37 wheel
+        python3 -m pip install pyqg==0.7.2 --no-build-isolation --use-deprecated=legacy-resolver
 
     - name: Lint with flake8
       run: |

README.md

@@ -20,50 +20,45 @@ Neural Networks,Volume 153, 530-552, ISSN 0893-6080, https://doi.org/10.1016/j.n
 
 ## Installation
 
+We recommend setting up a virtual environment using either [conda](https://docs.conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html) or [virtualenv](https://virtualenv.pypa.io/en/latest/user_guide.html).
+
 #### Clone Repo:
 
 ```bash
 git clone [email protected]:StevePny/DataAssimBench.git
 ```
 
-#### Set Up Conda Environment
-
-```bash
-cd DataAssimBench
-conda env create -f environment.yml
-conda activate dab
-```
-
 #### Install dabench
 ```bash
-pip install .
+cd ./DataAssimBench
+pip install -e ".[full]"
 ```
 
-#### Install dependencies (optional)
-The user may have to manually install:  
+This will create a full installation including the ability to access cloud data or interface with other packages such as qgs. Alternatively, for a minimal installation, run:
+
 ```bash
-conda install -c conda-forge jax  
-conda install -c conda-forge pyqg  
+pip install -e .
 ```
 
+
 ## Quick Start
 
 For more detailed examples, go to the [DataAssimBench-Examples](https://github.com/StevePny/DataAssimBench-Examples) repo.
 
 #### Importing data generators
 
 ```python
-from dabench import data
-help(data) # View data classes, e.g. data.Lorenz96
-help(data.Lorenz96) # Get more info about Lorenz96 class
+import dabench as dab
+help(dab.data) # View data classes, e.g. data.Lorenz96
+help(dab.data.Lorenz96) # Get more info about Lorenz96 class
 ```
 
 #### Generating data
 
 All of the data objects are set up with reasonable defaults. Generating data is as easy as:
 
 ```python
-l96_obj = data.Lorenz96() # Create data generator object
+l96_obj = dab.data.Lorenz96() # Create data generator object
 l96_obj.generate(n_steps=1000) # Generate Lorenz96 simulation data
 l96_obj.values # View the output values
 ```
@@ -76,7 +71,7 @@ All data objects are customizable.
 
 For data-generators (e.g. numerical models such as Lorenz63, Lorenz96, SQGTurb), this means you can change initial conditions, model parameters, timestep size, number of timesteps, etc.
 
-For data-downloaders (e.g. ENSOIDX, AWS, GCP), this means changing which variables you download, the lat/lon bounding box, the time period, etc.
+For data-downloaders (e.g. ENSOIDX, GCP), this means changing which variables you download, the lat/lon bounding box, the time period, etc.
 
 The recommended way of specifying options is to pass a keyword argument (kwargs) dictionary. The exact options vary between the different types of data objects, so be sure to check the specific documentation for your chosen generator/downloader more info.
 
@@ -86,18 +81,18 @@ The recommended way of specifying options is to pass a keyword argument (kwargs)
 l96_options = {'forcing_term': 7.5,
                'system_dim': 5,
                'delta_t': 0.05}
-l96_obj = data.Lorenz96(**l96_options) # Create data generator object
+l96_obj = dab.data.Lorenz96(**l96_options) # Create data generator object
 l96_obj.generate(n_steps=1000) # Generate Lorenz96 simulation data
 l96_obj.values # View the output values
 ```
 
-- For example, for the Amazon Web Services (AWS) ERA5 data-downloader, we can select our variables and time period like this:
+- For example, for the Google Cloud (GCP) ERA5 data-downloader, we can select our variables and time period like this:
 
 ```python
-aws_options = {'variables': ['air_pressure_at_mean_sea_level', 'sea_surface_temperature'],
-               'years': [1984, 1985]}
-aws_obj = data.AWS(**aws_options) # Create data generator object
-aws_obj.load() # Loads data. Can also use aws_obj.generate()
-aws_obj.values # View the output values
+gcp_options = {'variables': ['2m_temperature', 'sea_surface_temperature'],
+               'date_start': '2020-06-01'
+               'date_end': '2020-06-07'}
+gcp_obj = dab.data.GCP(**gcp_options) # Create data generator object
+gcp_obj.load() # Loads data. Can also use gcp_obj.generate()
+gcp_obj.values # View the output values
 ```
-

dabench/dacycler/_etkf.py

@@ -265,6 +265,9 @@ def cycle(self,
             (obs_vector.times > cur_time - analysis_window/2)
             # AND Less than end of window
             * (obs_vector.times < cur_time + analysis_window/2)
+            # AND not equal to end of window
+            * (1-jnp.isclose(obs_vector.times, cur_time + analysis_window/2,
+                             rtol=0))
             # OR Equal to start of window
             + jnp.isclose(obs_vector.times, cur_time - analysis_window/2,
                           rtol=0)

dabench/data/_data.py

@@ -275,6 +275,7 @@ def generate(self, n_steps=None, t_final=None, x0=None, M0=None,
     def _import_xarray_ds(self, ds, include_vars=None, exclude_vars=None,
                           years_select=None, dates_select=None,
                           lat_sorting=None):
+        ds = ds.as_numpy()
         if dates_select is not None:
             dates_filter_indices = ds.time.dt.date.isin(dates_select)
             # First check to make sure the dates exist in the object
@@ -313,8 +314,8 @@ def _import_xarray_ds(self, ds, include_vars=None, exclude_vars=None,
                           'stall, or crash.'.format(size_gb))
 
         # Get dims
-        dims = ds.dims
-        dims_names = list(ds.dims)
+        dims = ds.sizes
+        dims_names = list(ds.sizes)
 
         # Set times
         time_key = None
@@ -402,7 +403,7 @@ def _import_xarray_ds(self, ds, include_vars=None, exclude_vars=None,
                           )
 
         # Gather values and set dimensions
-        temp_values = np.moveaxis(np.array(ds.to_array()), 0, -1)
+        temp_values = np.moveaxis(ds.to_dataarray().to_numpy(), 0, -1)
         self.original_dim = temp_values.shape[1:]
         if self.original_dim[-1] == 1 and len(self.original_dim) > 2:
             self.original_dim = self.original_dim[:-1]
@@ -449,21 +450,13 @@ def load_netcdf(self, filepath=None, include_vars=None, exclude_vars=None,
         """
         if filepath is None:
             # Use importlib.resources to get the default netCDF from dabench
-            with resources.open_binary(
-                    _suppl_data, 'era5_japan_slp.nc') as nc_file:
-                with xr.open_dataset(nc_file, decode_coords='all') as ds:
-                    self._import_xarray_ds(
-                        ds, include_vars=include_vars,
-                        exclude_vars=exclude_vars,
-                        years_select=years_select, dates_select=dates_select,
-                        lat_sorting=lat_sorting)
-        else:
-            with xr.open_dataset(filepath, decode_coords='all') as ds:
-                self._import_xarray_ds(
-                    ds, include_vars=include_vars,
-                    exclude_vars=exclude_vars,
-                    years_select=years_select, dates_select=dates_select,
-                    lat_sorting=lat_sorting)
+            filepath = resources.files(_suppl_data).joinpath('era5_japan_slp.nc')
+        with xr.open_dataset(filepath, decode_coords='all') as ds:
+            self._import_xarray_ds(
+                ds, include_vars=include_vars,
+                exclude_vars=exclude_vars,
+                years_select=years_select, dates_select=dates_select,
+                lat_sorting=lat_sorting)
 
     def save_netcdf(self, filename):
         """Saves values in values attribute to netCDF file

dabench/data/enso_indices.py

@@ -275,7 +275,7 @@ def _combine_vals_years(self, all_vals, all_years):
             all_vals[v] = all_vals[v][np.sort(indices)]
             all_years[v] = all_years[v][np.sort(indices)]
             # Append common_vals
-            common_vals.append(all_vals[v][jnp.in1d(all_years[v],
+            common_vals.append(all_vals[v][jnp.isin(all_years[v],
                                                     common_years)])
         for f in common_vals:
             print(f.shape)

dabench/data/gcp.py

@@ -129,8 +129,8 @@ def _load_gcp_era5(self):
         if self.min_lat is not None and self.max_lat is not None:
             # Subset by lat boundaries
             ds = ds.where(
-                (ds.latitude < self.max_lat) &
-                (ds.latitude > self.min_lat),
+                ((ds.latitude < self.max_lat) &
+                 (ds.latitude > self.min_lat)).compute(),
                 drop=True)
         if self.min_lon is not None and self.max_lon is not None:
             # Convert west longs to degrees east
@@ -142,8 +142,8 @@ def _load_gcp_era5(self):
                 subset_max_lon += 360
             # Subset by lon boundaries
             ds = ds.where(
-                (ds.longitude < subset_max_lon) &
-                (ds.longitude > subset_min_lon),
+                ((ds.longitude < subset_max_lon) &
+                 (ds.longitude > subset_min_lon)).compute(),
                 drop=True)
 
         self._import_xarray_ds(ds)

dabench/data/sqgturb.py

@@ -36,15 +36,14 @@
 import jax
 import jax.numpy as jnp
 from jax.numpy.fft import rfft2, irfft2
-from jax.config import config
 from functools import partial
 from importlib import resources
 
 from dabench.data import _data
 from dabench import _suppl_data
 
 # Set to enable 64bit floats in Jax
-config.update('jax_enable_x64', True)
+jax.config.update('jax_enable_x64', True)
 
 
 class SQGTurb(_data.Data):

dabench/obsop/_obsop.py

@@ -29,7 +29,7 @@ def __init__(self,
             self.h = self._index_state_vec
         else:
             # Check custom h
-            custom_args = inspect.getargspec(h).args
+            custom_args = inspect.signature(h).parameters
             if 'state_vec' not in custom_args or 'obs_vec' not in custom_args:
                 raise ValueError('User-specified h does not accept the '
                                  'required args for h: "state_vec" and '