Below are some code bases that might be useful for people to get started.
Algorithms¶
Differentiable Ocean Models¶
jaxsw. This library tries to . It uses components from other libraries like equinox for PyTrees, FiniteDiffX for finite differences, kernex for interpretable convolution operators, and diffrax for time steppers.
veros. This library is a fully-fledged Ocean GCM.
4DVarNet¶
This is a set of code bases that were used to design an end-to-end learning system for SSH interpolation.
4DVarNet-Core [PyTorch]
4DVarNet-Starter [PyTorch]
Core Components¶
PDE Elements¶
FiniteDiffX - Finite Differences with JAX.
kernex - stencil operations with JAX.
diffrax - ODE time steppers with JAX.
Optimization¶
JaxOpt - argmin differentiation with JAX.
Lineax - linear solvers with JAX.
Simple Ocean Models¶
- pyQG - Abernathey et al., 2022
- jax-cfd - Kochkov et al. (2021)Dresdner et al. (2022)
GCMs
- veros - Häfner et al., 2018
- Oceananigans - Ramadhan et al., 2020
Data Assimilation¶
- Abernathey, R., Rochanotes, Ross, A., Jansen, M., Ziwei Li, Poulin, F. J., Constantinou, N. C., Anirban Sinha, Dhruv Balwada, SalahKouhen, Jones, S., Rocha, C. B., Wolfe, C. L. P., Chuizheng Meng, Van Kemenade, H., Bourbeau, J., Penn, J., Busecke, J., Bueti, M., & , Tobias. (2022). pyqg/pyqg: v0.7.2. Zenodo. 10.5281/ZENODO.6563667
- Kochkov, D., Smith, J. A., Alieva, A., Wang, Q., Brenner, M. P., & Hoyer, S. (2021). Machine learning–accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences, 118(21). 10.1073/pnas.2101784118
- Dresdner, G., Kochkov, D., Norgaard, P., Zepeda-Núñez, L., Smith, J. A., Brenner, M. P., & Hoyer, S. (2022). Learning to correct spectral methods for simulating turbulent flows. arXiv. 10.48550/ARXIV.2207.00556
- Häfner, D., Jacobsen, R. L., Eden, C., Kristensen, M. R. B., Jochum, M., Nuterman, R., & Vinter, B. (2018). Veros v0.1 – a fast and versatile ocean simulator in pure Python. Geoscientific Model Development, 11(8), 3299–3312. 10.5194/gmd-11-3299-2018
- Ramadhan, A., Wagner, G., Hill, C., Campin, J.-M., Churavy, V., Besard, T., Souza, A., Edelman, A., Ferrari, R., & Marshall, J. (2020). Oceananigans.jl: Fast and friendly geophysical fluid dynamics on GPUs. Journal of Open Source Software, 5(53), 2018. 10.21105/joss.02018