initial commit

.gitignore

@@ -0,0 +1,16 @@
+
+# vim tmp
+*.swp
+
+# python
+__pycache__/
+.ipynb_checkpoints/
+*.egg-info
+dist/
+
+# sphinx doc
+docs/_build/
+
+# personal tmp
+tmp/
+docs/pages/notebooks/pyshqg

.readthedocs.yaml

@@ -0,0 +1,13 @@
+version: "2"
+
+build:
+  os: "ubuntu-22.04"
+  tools:
+    python: "3.10"
+
+python:
+  install:
+    - requirements: docs/requirements.txt
+
+sphinx:
+  configuration: docs/conf.py

CHANGELOG.md

@@ -0,0 +1,10 @@
+# Version history
+
+## Version 1.0.0 (2024-03-12)
+
+- First release version.
+
+## Version 1.0.1 (2024-03-12)
+
+- Fixed test datasets.
+

CITATION.cff

@@ -0,0 +1,15 @@
+cff-version: 1.2.0
+message: "If you use this software, please cite it as below."
+authors:
+- family-names: "Farchi"
+  given-names: "Alban"
+  orcid: "https://orcid.org/0000-0002-4162-8289"
+- family-names: "Malartic"
+  given-names: "Quentin"
+- family-names: "Bocquet"
+  given-names: "Marc"
+title: "The Python Spherical Harmonics Quasi-Geostrophic (pyshqg) model"
+version: 0.1.0
+doi: 
+date-released: 2023-12-04
+url:

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Alban Farchi, Quentin Malartic, Marc Bocquet, and Fabio D'Andrea
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

MANIFEST.in

@@ -0,0 +1,61 @@
+include pyshqg/data/data_t63.zarr/.zattrs
+include pyshqg/data/data_t63.zarr/.zgroup
+include pyshqg/data/data_t63.zarr/.zmetadata
+include pyshqg/data/data_t63.zarr/forcing/0.0.0
+include pyshqg/data/data_t63.zarr/initial_q/0.0.0
+include pyshqg/data/data_t63.zarr/land_sea_mask/0.0
+include pyshqg/data/data_t63.zarr/lat/0
+include pyshqg/data/data_t63.zarr/lon/0
+include pyshqg/data/data_t63.zarr/orography/0.0
+include pyshqg/data/test_t21_t31.zarr/.zattrs
+include pyshqg/data/test_t21_t31.zarr/.zgroup
+include pyshqg/data/test_t21_t31.zarr/.zmetadata
+include pyshqg/data/test_t21_t31.zarr/dissip_ekman/0.0.0.0
+include pyshqg/data/test_t21_t31.zarr/dpsi_dphi/0.0.0.0
+include pyshqg/data/test_t21_t31.zarr/dpsi_dtheta/0.0.0.0
+include pyshqg/data/test_t21_t31.zarr/dq_dphi/0.0.0.0
+include pyshqg/data/test_t21_t31.zarr/dq_dtheta/0.0.0.0
+include pyshqg/data/test_t21_t31.zarr/forcing/0.0.0
+include pyshqg/data/test_t21_t31.zarr/grid_dq_dt/0.0.0.0
+include pyshqg/data/test_t21_t31.zarr/grid_to_spec_dq_dt/0.0.0.0.0
+include pyshqg/data/test_t21_t31.zarr/jacobian/0.0.0.0
+include pyshqg/data/test_t21_t31.zarr/lat/0
+include pyshqg/data/test_t21_t31.zarr/lon/0
+include pyshqg/data/test_t21_t31.zarr/spec_abm/0.0.0.0.0
+include pyshqg/data/test_t21_t31.zarr/spec_dissip_sel/0.0.0.0.0
+include pyshqg/data/test_t21_t31.zarr/spec_dissip_therm/0.0.0.0.0
+include pyshqg/data/test_t21_t31.zarr/spec_dq_dt/0.0.0.0.0
+include pyshqg/data/test_t21_t31.zarr/spec_ee/0.0.0.0.0
+include pyshqg/data/test_t21_t31.zarr/spec_psi/0.0.0.0.0
+include pyshqg/data/test_t21_t31.zarr/spec_q/0.0.0.0.0
+include pyshqg/data/test_t21_t31.zarr/spec_rk2/0.0.0.0.0
+include pyshqg/data/test_t21_t31.zarr/spec_rk4/0.0.0.0.0
+include pyshqg/data/test_t21_t31.zarr/spec_tendencies/0.0.0.0.0
+include pyshqg/data/test_t21_t31.zarr/spec_total_q/0.0.0.0.0
+include pyshqg/data/test_t21_t31.zarr/zeta/0.0.0
+include pyshqg/data/test_t42_t63.zarr/.zattrs
+include pyshqg/data/test_t42_t63.zarr/.zgroup
+include pyshqg/data/test_t42_t63.zarr/.zmetadata
+include pyshqg/data/test_t42_t63.zarr/dissip_ekman/0.0.0.0
+include pyshqg/data/test_t42_t63.zarr/dpsi_dphi/0.0.0.0
+include pyshqg/data/test_t42_t63.zarr/dpsi_dtheta/0.0.0.0
+include pyshqg/data/test_t42_t63.zarr/dq_dphi/0.0.0.0
+include pyshqg/data/test_t42_t63.zarr/dq_dtheta/0.0.0.0
+include pyshqg/data/test_t42_t63.zarr/forcing/0.0.0
+include pyshqg/data/test_t42_t63.zarr/grid_dq_dt/0.0.0.0
+include pyshqg/data/test_t42_t63.zarr/grid_to_spec_dq_dt/0.0.0.0.0
+include pyshqg/data/test_t42_t63.zarr/jacobian/0.0.0.0
+include pyshqg/data/test_t42_t63.zarr/lat/0
+include pyshqg/data/test_t42_t63.zarr/lon/0
+include pyshqg/data/test_t42_t63.zarr/spec_abm/0.0.0.0.0
+include pyshqg/data/test_t42_t63.zarr/spec_dissip_sel/0.0.0.0.0
+include pyshqg/data/test_t42_t63.zarr/spec_dissip_therm/0.0.0.0.0
+include pyshqg/data/test_t42_t63.zarr/spec_dq_dt/0.0.0.0.0
+include pyshqg/data/test_t42_t63.zarr/spec_ee/0.0.0.0.0
+include pyshqg/data/test_t42_t63.zarr/spec_psi/0.0.0.0.0
+include pyshqg/data/test_t42_t63.zarr/spec_q/0.0.0.0.0
+include pyshqg/data/test_t42_t63.zarr/spec_rk2/0.0.0.0.0
+include pyshqg/data/test_t42_t63.zarr/spec_rk4/0.0.0.0.0
+include pyshqg/data/test_t42_t63.zarr/spec_tendencies/0.0.0.0.0
+include pyshqg/data/test_t42_t63.zarr/spec_total_q/0.0.0.0.0
+include pyshqg/data/test_t42_t63.zarr/zeta/0.0.0

README.md

@@ -0,0 +1,111 @@
+
+![Animated example](http://cerea.enpc.fr/HomePages/farchia/pyshqg-map.gif)
+
+# pyshqg
+
+[![Documentation Status](https://readthedocs.org/projects/pyshqg/badge/?version=latest)](https://pyshqg.readthedocs.io/en/latest/?badge=latest)
+[![PyPI version](https://badge.fury.io/py/pyshqg.svg)](https://badge.fury.io/py/pyshqg)
+
+`pyshqg` is a python solver for the [Marshal and Molteni (1993) quasi-geostrophic (QG) model](https://doi.org/10.1175/1520-0469(1993)050%3C1792:TADUOP%3E2.0.CO;2).
+QG models express the conservation of potential vorticity over time and are
+meant to describe the large-scale circulation of the atmosphere under specific hypotheses.
+This QG model is very special, because it is expressed in spherical harmonics and
+because it encapsulates complex physical processes.
+
+- [Documentation](https://pyshqg.readthedocs.io)
+- [Source code](https://github.com/cerea-daml/pyshqg)
+- [Issue tracker](https://github.com/cerea-daml/pyshqg/issues)
+
+## Installation
+
+Install using pip:
+
+    $ pip install pyshqg
+
+More details can be found on this [page](https://pyshqg.readthedocs.io/en/latest/pages/installation.html).
+
+## Usage
+
+Here is a sneak peak at how to use the package.
+First import all the machinery.
+
+    >>> import numpy as np
+    >>>
+    >>> from pyshqg.backend.numpy_backend import NumpyBackend as Backend
+    >>> from pyshqg.preprocessing.reference_data import load_test_data
+    >>> from pyshqg.core.constructors import construct_model, construct_integrator
+
+Load the data and create the model and integrator.
+
+    >>> ds_test = load_test_data(internal_truncature=21, grid_truncature=31)
+    >>> backend = Backend(floatx='float64')
+    >>> model = construct_model(backend, ds_test.config)
+    >>> rk4 = construct_integrator(ds_test.config['rk4_integration'], model)
+
+Perform one integration step and test against the test data.
+
+    >>> state = model.model_state(ds_test.spec_q.to_numpy())
+    >>> fwd_state = rk4.forward(state)
+    >>>
+    >>> def rms(x):
+    ...     return np.sqrt(np.mean(np.square(x)))
+    >>>
+    >>> def rmse(x, y):
+    ...     return rms(x-y)
+    >>>
+    >>> rms_ref = rms(ds_test.spec_rk4.to_numpy())
+    >>> rms_diff = rmse(
+    ...     backend.to_numpy(fwd_state['spec_q']),
+    ...     ds_test.spec_rk4.to_numpy(),
+    ... )
+    >>> assert rms_diff < 1e-6 * rms_ref
+
+Run the model for a full trajectory and convert the output to an ``xarray.Dataset``.
+
+    >>> state = model.model_state(ds_test.spec_q.to_numpy())
+    >>> trajectory = rk4.run(
+    ...     state,
+    ...     t_start=0,
+    ...     num_snapshots=8,
+    ...     num_steps_per_snapshot=10,
+    ...     variables=('q', 'psi'),
+    ... )
+    >>> trajectory.to_xarray()
+    <xarray.Dataset> Size: 14MB
+    Dimensions:  (time: 9, batch: 16, level: 3, lat: 32, lon: 64)
+    Coordinates:
+      * time     (time) int64 72B 0 18000 36000 54000 ... 90000 108000 126000 144000
+      * lat      (lat) float64 256B -85.76 -80.27 -74.74 ... 74.74 80.27 85.76
+      * lon      (lon) float64 512B 180.0 185.6 191.2 196.9 ... 163.1 168.8 174.4
+    Dimensions without coordinates: batch, level
+    Data variables:
+        q        (time, batch, level, lat, lon) float64 7MB -0.000298 ... -0.0001269
+        psi      (time, batch, level, lat, lon) float64 7MB 1.445e+08 ... -1.38e+07
+
+More details can be found on this [page](https://pyshqg.readthedocs.io/en/latest/pages/examples.html).
+
+## Aknowledgements
+
+This python package is based on an original implementation of the model
+in Fortran written at LMD by XXX.
+
+## Todo-list
+
+- make higher resolution gif animation
+- make more example notebooks (see list below)
+- fill in the documentation sections
+- fill in docstrings, in particular with the convention used in the docs
+- publish first version on pypi and conda-forge
+- write pytorch/jax backends
+- include unit tests in github CI
+- apply linting and include it in github CI
+
+Potential example notebooks:
+
+- one notebook for each term of the QG model
+- how to compute the forecast skill T15 vs T31 (as illustration)
+- how to use tf backend to compute the grad
+- gradient test for tf backend
+- adjoint test for tf backend
+- how to define a NN-based model error correction for the T21 and train it with dummy data
+

docs/Makefile

@@ -0,0 +1,20 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = .
+BUILDDIR      = _build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

docs/conf.py

@@ -0,0 +1,108 @@
+# Configuration file for the Sphinx documentation builder.
+#
+# This file only contains a selection of the most common options. For a full
+# list see the documentation:
+# https://www.sphinx-doc.org/en/master/usage/configuration.html
+
+# -- Path setup --------------------------------------------------------------
+
+# If extensions (or modules to document with autodoc) are in another directory,
+# add these directories to sys.path here. If the directory is relative to the
+# documentation root, use os.path.abspath to make it absolute, like shown here.
+
+import os
+import sys
+import tomli
+import datetime
+
+#import pyshqg
+sys.path.insert(0, os.path.abspath('..'))
+
+
+# -- Project information -----------------------------------------------------
+
+with open('../pyproject.toml', mode='rb') as pyproject:
+    pkg_meta = tomli.load(pyproject)['project']
+
+project = str(pkg_meta['name'])
+copyright = str(datetime.datetime.now().year) + ", CEREA-DAML Team"
+author = 'CEREA-DAML Team'
+version = str(pkg_meta['version'])
+release = version
+
+
+# -- General configuration ---------------------------------------------------
+
+# Add any Sphinx extension module names here, as strings. They can be
+# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
+# ones.
+extensions = [
+    'sphinx.ext.autodoc',
+    'sphinx.ext.mathjax',
+    'sphinx.ext.napoleon',
+    'sphinx.ext.coverage',
+    # used to include .md files
+    'm2r2',
+    # theme
+    'sphinx_rtd_theme',
+    # no have notebooks
+    'nbsphinx',
+    'sphinx_gallery.load_style',
+]
+
+autodoc_mock_imports = [
+    'numpy',
+    'xarray',
+    'pyshtools',
+    'tensorflow',
+    'tqdm',
+    'scipy',
+]
+
+autoclass_content = 'class'
+autodoc_member_order = 'bysource'
+autodoc_default_flags = {
+    'members': '',
+    'undoc-members': 'code,error_template',
+    'exclude-members': '__dict__,__weakref__',
+}
+source_suffix = ['.rst', '.md']
+language = 'en'
+pygments_style = 'sphinx'
+add_module_names = True
+autodoc_default_options = {
+    'show-inheritance': False,
+}
+coverage_show_missing_items = True
+# Napoleon settings
+napoleon_google_docstring = True
+napoleon_numpy_docstring = True
+napoleon_include_init_with_doc = True
+napoleon_include_private_with_doc = True
+napoleon_include_special_with_doc = True
+napoleon_use_admonition_for_examples = False
+napoleon_use_admonition_for_notes = False
+napoleon_use_admonition_for_references = False
+napoleon_use_ivar = False
+napoleon_use_param = True
+napoleon_use_rtype = True
+
+# Add any paths that contain templates here, relative to this directory.
+templates_path = ['_templates']
+
+# List of patterns, relative to source directory, that match files and
+# directories to ignore when looking for source files.
+# This pattern also affects html_static_path and html_extra_path.
+exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']
+
+# -- Options for HTML output -------------------------------------------------
+
+# The theme to use for HTML and HTML Help pages.  See the documentation for
+# a list of builtin themes.
+#
+html_theme = 'sphinx_rtd_theme'
+
+# Add any paths that contain custom static files (such as style sheets) here,
+# relative to this directory. They are copied after the builtin static files,
+# so a file named "default.css" will overwrite the builtin "default.css".
+html_static_path = ['_static']