diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..a54d1be --- /dev/null +++ b/.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 }} diff --git a/README.md b/README.md index ad9e049..93dd205 100644 --- a/README.md +++ b/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. diff --git a/pynanigans/grids.py b/pynanigans/grids.py index b845536..62cf597 100644 --- a/pynanigans/grids.py +++ b/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) diff --git a/pynanigans/utils.py b/pynanigans/utils.py index ec1ad38..f962c05 100644 --- a/pynanigans/utils.py +++ b/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) diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..f6e9afa --- /dev/null +++ b/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", + ], +) diff --git a/tests/test_grids.py b/tests/test_grids.py new file mode 100644 index 0000000..3103580 --- /dev/null +++ b/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" diff --git a/tests/test_pnplot.py b/tests/test_pnplot.py new file mode 100644 index 0000000..9235ca2 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 0000000..4b9a22a --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,80 @@ +import pytest +import xarray as xr +import numpy as np +from pynanigans.utils import biject, normalize_time_by, downsample, pnchunk + +def test_biject(): + # Create a test dataset + data = np.random.rand(10, 10, 10) + dims = ['xC', 'yC', 'zC'] + coords = { + 'xC': np.linspace(0, 1, 10), + 'yC': np.linspace(0, 1, 10), + 'zC': np.linspace(0, 1, 10) + } + ds = xr.Dataset( + data_vars={'u': (dims, data)}, + coords=coords + ) + + # Test bijection + result = biject(ds.u) + assert 'x' in result.dims + assert 'y' in result.dims + assert 'z' in result.dims + assert 'xC' not in result.dims + assert 'yC' not in result.dims + assert 'zC' not in result.dims + +def test_normalize_time_by(): + # Create a test dataset with time + data = np.random.rand(10) + time = np.array([np.timedelta64(i, 'ns') for i in range(10)]) + ds = xr.Dataset( + data_vars={'u': ('time', data)}, + coords={'time': time} + ) + + # Test normalization + result = normalize_time_by(ds.u, seconds=1) + assert result.time.dtype == 'float64' + assert result.time.attrs['units'] == 'seconds' + +def test_downsample(): + # Create a test dataset + data = np.random.rand(100, 100) + dims = ['xC', 'yC'] + coords = { + 'xC': np.linspace(0, 1, 100), + 'yC': np.linspace(0, 1, 100) + } + ds = xr.Dataset( + data_vars={'u': (dims, data)}, + coords=coords + ) + + # Test downsampling + result = downsample(ds.u, xC=50, yC=50) + assert len(result.xC) == 50 + assert len(result.yC) == 50 + +def test_pnchunk(): + # Create a test dataset with time + data = np.random.rand(100, 10, 10, 10) + time = np.array([np.timedelta64(i, 'ns') for i in range(100)]) + dims = ['time', 'xC', 'yC', 'zC'] + coords = { + 'time': time, + 'xC': np.linspace(0, 1, 10), + 'yC': np.linspace(0, 1, 10), + 'zC': np.linspace(0, 1, 10) + } + ds = xr.Dataset( + data_vars={'u': (dims, data)}, + coords=coords + ) + + # Test chunking + result = pnchunk(ds.u, maxsize_4d=1000) + result = pnchunk(ds, maxsize_4d=1000) + assert any(result.chunks)