Add tests

.github/workflows/tests.yml

@@ -0,0 +1,39 @@
+name: Tests
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: [3.9, 3.11, 3.12, 3.13]
+
+    steps:
+    - uses: actions/checkout@v2
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v2
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install pytest pytest-cov
+        pip install -e .
+
+    - name: Run tests
+      run: |
+        pytest tests/ --cov=pynanigans --cov-report=xml
+
+    - name: Upload coverage to Codecov
+      uses: codecov/codecov-action@v5
+      with:
+        files: ./coverage.xml
+        fail_ci_if_error: false
+        token: ${{ secrets.CODECOV_TOKEN }}

README.md

@@ -1,10 +1,80 @@
 # pynanigans
+
 Python scripts for Oceananigans.jl NetCDF output
 
 ## Installation
 
-- Clone this repo somewhere
-- Install with pip:
-  - `pip install /path/to/cloned/repo`
+### Using pip
+
+You can install pynanigans directly from the repository:
+
+```bash
+pip install git+https://github.com/yourusername/pynanigans.git
+```
+
+Or install from a local clone:
+
+```bash
+git clone https://github.com/yourusername/pynanigans.git
+cd pynanigans
+pip install -e .
+```
+
+### Using conda
+
+If you prefer using conda, you can create an environment with all dependencies:
+
+```bash
+# Clone the repository
+git clone https://github.com/yourusername/pynanigans.git
+cd pynanigans
+
+# Create and activate conda environment
+conda env create -f environment.yml
+conda activate p39
+
+# Install the package
+pip install -e .
+```
+
+## Dependencies
+
+The package requires:
+- Python >= 3.9
+- numpy
+- xarray
+- xgcm
+- matplotlib
+
+## Development
+
+To set up a development environment:
+
+1. Clone the repository:
+```bash
+git clone https://github.com/yourusername/pynanigans.git
+cd pynanigans
+```
+
+2. Create a virtual environment and install dependencies:
+```bash
+python -m venv venv
+source venv/bin/activate  # On Windows: venv\Scripts\activate
+pip install -e ".[dev]"
+```
+
+3. Install development dependencies:
+```bash
+pip install pytest pytest-cov
+```
+
+4. Run tests:
+```bash
+pytest tests/
+```
+
+## License
+
+See the [LICENSE](LICENSE) file for details.
 
 

pynanigans/grids.py

@@ -1,6 +1,6 @@
 
 def get_coords(ds, topology="PPN",):
-    """ 
+    """
     Constructs the coords dict for ds to be passed to xgcm.Grid
     Flat dimensions (F) are treated the same as Periodic ones (P)
     """
@@ -9,7 +9,7 @@ def get_coords(ds, topology="PPN",):
     per = { dim : dict(left=f"{dim}F", center=f"{dim}C") for dim in "xyz" }
     nper = { dim : dict(outer=f"{dim}F", center=f"{dim}C") for dim in "xyz" }
     coords = { dim : per[dim] if top in "FP" else nper[dim] for dim, top in zip("xyz", topology) }
-    
+
     return coords
 
 
@@ -38,7 +38,7 @@ def get_distances(ds, dim="x", topology="P"):
 
 
 def get_metrics(ds, topology="PPN"):
-    """ 
+    """
     Constructs the metric dict for ds.
     (Not sure if the metrics are correct at the boundary points
     """
@@ -67,7 +67,7 @@ def get_grid(ds, coords=None, metrics=None, topology="PPN", **kwargs):
         metrics = get_metrics(ds, topology=topology)
 
     periodic = [ dim for (dim, top) in zip("xyz", topology) if top in "PF" ]
-    return xg.Grid(ds, coords=coords, metrics=metrics, 
+    return xg.Grid(ds, coords=coords, metrics=metrics,
                    periodic=periodic, **kwargs)
 
 

pynanigans/utils.py

@@ -14,10 +14,10 @@ def biject(darray, *args, surjection=surjection):
     """
     Renames darray so that the dimension names actually correspond to the physical
     dimensions, instead of being the name of the grid meshes in Oceananigans.
-    If `*args` is provided, only those dimensions will be renamed. If not, `x`, `y` 
+    If `*args` is provided, only those dimensions will be renamed. If not, `x`, `y`
     and `z` will be automatically renamed.
 
-    This makes calling functions easier as instead of calling `darray.u.plot(x='xF', y='zC')`, 
+    This makes calling functions easier as instead of calling `darray.u.plot(x='xF', y='zC')`,
     you can call `darray.pnplot(x='x', y='z')`
     """
     da_dims = darray.dims
@@ -39,7 +39,7 @@ def regular_indef_integrate(f, dim):
 
 
 def normalize_time_by(darray, seconds=1, new_units="seconds"):
-    """ 
+    """
     Converts the time dimension (a timedelta[ns] object by default) into a np.float64
     object while normalizing it by number of seconds `seconds`.
     """
@@ -72,13 +72,15 @@ def downsample(darray, round_func=round, **dim_limits):
 
 
 def pnchunk(darray, maxsize_4d=1000**2, sample_var="u", round_func=round, **kwargs):
-    """ Chunk `darray` in time while keeping each chunk's size roughly 
+    """ Chunk `darray` in time while keeping each chunk's size roughly
     around `maxsize_4d`. The default `maxsize_4d=1000**2` comes from
     xarray's rule of thumb for chunking:
     http://xarray.pydata.org/en/stable/dask.html#chunking-and-performance
     """
-    chunk_number = darray[sample_var].size / maxsize_4d
-    chunk_size = int(round_func(len(darray[sample_var].time) / chunk_number))
+    if type(darray) == xr.Dataset:
+        darray = darray[sample_var]
+    chunk_number = darray.size / maxsize_4d
+    chunk_size = int(round_func(len(darray.time) / chunk_number))
     return darray.chunk(dict(time=chunk_size))
 xr.DataArray.pnchunk = pnchunk
 xr.Dataset.pnchunk = pnchunk
@@ -119,7 +121,7 @@ def pn{funcname}(darray, dim=None, surjection=surjection, **kwargs):
     exec(funcdef)
 
 
-def open_simulation(fname, 
+def open_simulation(fname,
                     grid_kwargs=dict(),
                     load=False,
                     squeeze=True,
@@ -131,7 +133,7 @@ def open_simulation(fname,
 
     `kwargs` are passed to `xarray.open_dataset()` and `grid_kwargs` are passed to `pynanigans.get_grid()`.
     """
-    
+
     #++++ Open dataset and create grid before squeezing
     if load:
         ds = xr.load_dataset(fname, **kwargs)

setup.py

@@ -0,0 +1,26 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="pynanigans",
+    version="0.1.0",
+    packages=find_packages(),
+    install_requires=[
+        "numpy",
+        "xarray",
+        "xgcm",
+        "matplotlib",
+    ],
+    python_requires=">=3.9",
+    author="Tomas Chor",
+    author_email="contact@tomaschor.xyz",
+    description="A Python package for working with Oceananigans data",
+    long_description=open("README.md").read(),
+    long_description_content_type="text/markdown",
+    url="https://github.com/yourusername/pynanigans",
+    classifiers=[
+        "Development Status :: 3 - Alpha",
+        "Intended Audience :: Science/Research",
+        "License :: OSI Approved :: MIT License",
+        "Programming Language :: Python",
+    ],
+)

tests/test_grids.py

@@ -0,0 +1,18 @@
+import pytest
+import xarray as xr
+import numpy as np
+from pynanigans.grids import get_coords, get_metrics, get_grid
+
+def test_get_coords():
+    # Test periodic coordinates
+    coords = get_coords(None, topology="PPP")
+    assert "x" in coords
+    assert "y" in coords
+    assert "z" in coords
+    assert coords["x"]["left"] == "xF"
+    assert coords["x"]["center"] == "xC"
+
+    # Test non-periodic coordinates
+    coords = get_coords(None, topology="NNN")
+    assert coords["x"]["outer"] == "xF"
+    assert coords["x"]["center"] == "xC"

tests/test_pnplot.py

@@ -0,0 +1,89 @@
+import pytest
+import xarray as xr
+import numpy as np
+from pynanigans.pnplot import pnplot, _imshow, _pcolormesh, _contour, _contourf
+
+def test_pnplot():
+    # Create a test dataset
+    data = np.random.rand(10, 10)
+    dims = ['xC', 'yC']
+    coords = {
+        'xC': np.linspace(0, 1, 10),
+        'yC': np.linspace(0, 1, 10)
+    }
+    ds = xr.Dataset(
+        data_vars={'u': (dims, data)},
+        coords=coords
+    )
+
+    # Test plotting
+    plot = pnplot(ds.u, x='x', y='y')
+    assert plot is not None
+
+def test_imshow():
+    # Create a test dataset
+    data = np.random.rand(10, 10)
+    dims = ['xC', 'yC']
+    coords = {
+        'xC': np.linspace(0, 1, 10),
+        'yC': np.linspace(0, 1, 10)
+    }
+    ds = xr.Dataset(
+        data_vars={'u': (dims, data)},
+        coords=coords
+    )
+
+    # Test imshow
+    plot = _imshow(ds.u, x='x', y='y')
+    assert plot is not None
+
+def test_pcolormesh():
+    # Create a test dataset
+    data = np.random.rand(10, 10)
+    dims = ['xC', 'yC']
+    coords = {
+        'xC': np.linspace(0, 1, 10),
+        'yC': np.linspace(0, 1, 10)
+    }
+    ds = xr.Dataset(
+        data_vars={'u': (dims, data)},
+        coords=coords
+    )
+
+    # Test pcolormesh
+    plot = _pcolormesh(ds.u, x='x', y='y')
+    assert plot is not None
+
+def test_contour():
+    # Create a test dataset
+    data = np.random.rand(10, 10)
+    dims = ['xC', 'yC']
+    coords = {
+        'xC': np.linspace(0, 1, 10),
+        'yC': np.linspace(0, 1, 10)
+    }
+    ds = xr.Dataset(
+        data_vars={'u': (dims, data)},
+        coords=coords
+    )
+
+    # Test contour
+    plot = _contour(ds.u, x='x', y='y')
+    assert plot is not None
+
+def test_contourf():
+    # Create a test dataset
+    data = np.random.rand(10, 10)
+    dims = ['xC', 'yC']
+    coords = {
+        'xC': np.linspace(0, 1, 10),
+        'yC': np.linspace(0, 1, 10)
+    }
+    ds = xr.Dataset(
+        data_vars={'u': (dims, data)},
+        coords=coords
+    )
+
+    # Test contourf
+    plot = _contourf(ds.u, x='x', y='y')
+    assert plot is not None