Updated docs.

_toc.yml

@@ -18,15 +18,15 @@
 format: jb-book
 root: content/overview
 chapters:
-  - title: Package Components
+  - title: Tutorials
     sections:
       - file: content/components/anatomy_pde
       - file: content/components/design
+      - file: content/components/grid_operations
       - file: content/components/boundaries
-  - title: Geometry
+  - title: Special Cases
     sections:
       - file: content/geometry/latlon_deltas
-      - file: content/geometry/staggering
   - title: Lorenz
     sections:
       - file: content/lorenz/lorenz
@@ -45,13 +45,13 @@ chapters:
       - file: content/12_steps/2.4_burgers_2d
       - file: content/12_steps/3.1_laplace
       - file: content/12_steps/3.2_poisson
-  - title: Quasi-Geostrophic
+  - title: Quasi-Geostrophic Model
     sections:
       - file: content/qg/qg
       - file: content/qg/ssh_free_run
       - file: content/qg/multilayer_qg
       - file: content/qg/qg_inversion
-  - title: Shallow Water
+  - title: Shallow Water Model
     sections:
       - file: content/sw/sw
       - file: content/sw/sw_linear_api1

content/geometry/staggering.ipynb → content/components/grid_operations.ipynb

@@ -4,7 +4,31 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Staggered Grids"
+    "---\n",
+    "title: Grid Operations\n",
+    "subject: Jax Approximate Ocean Models\n",
+    "subtitle: How can I use the Grid Operations to deal with staggered grids?\n",
+    "short_title: Grid Operations\n",
+    "authors:\n",
+    "  - name: J. Emmanuel Johnson\n",
+    "    affiliations:\n",
+    "      - CNRS\n",
+    "      - MEOM\n",
+    "    orcid: 0000-0002-6739-0053\n",
+    "    email: jemanjohnson34@gmail.com\n",
+    "  - name: Takaya Uchida\n",
+    "    affiliations:\n",
+    "      - FSU\n",
+    "    orcid: https://orcid.org/0000-0002-8654-6009\n",
+    "    email: tuchida@fsu.edu\n",
+    "license: CC-BY-4.0\n",
+    "keywords: jax, shallow water model, differentiable\n",
+    "abbreviations:\n",
+    "    SW: Shallow Water\n",
+    "    QG: Quasi-Geostrophic\n",
+    "    PDE: Partial Differential Equation\n",
+    "    RHS: Right Hand Side\n",
+    "---"
    ]
   },
   {
@@ -46,7 +70,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 1-Dimensional"
+    "## 1D Arakawa C-Grid"
    ]
   },
   {
@@ -395,7 +419,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Arakawa C-Grid"
+    "## 2D Arakawa C-Grid"
    ]
   },
   {

content/overview.md

@@ -1,55 +1,190 @@
-# Home Page
+---
+title: Home Page
+subject: Jax Approximate Ocean Models
+# subtitle: How can I estimate the state AND the parameters?
+short_title: Start Here
+authors:
+  - name: J. Emmanuel Johnson
+    affiliations:
+      - CNRS
+      - MEOM
+    orcid: 0000-0002-6739-0053
+    email: jemanjohnson34@gmail.com
+  - name: Takaya Uchida
+    affiliations:
+      - FSU
+    orcid: https://orcid.org/0000-0002-8654-6009
+    email: tuchida@fsu.edu
+license: CC-BY-4.0
+keywords: jax, shallow water model, differentiable
+abbreviations:
+    SW: Shallow Water
+    SWM: Shallow Water Model
+    QG: Quasi-Geostrophic
+    PDE: Partial Differential Equation
+    RHS: Right Hand Side
+---
 
 > A simple package for doing simple differential ocean model approximations in `JAX`.
 
 
 ---
+## Motivation
+
+
+
+:::{seealso} Other Packages
+:class: dropdown
+
+[**PhiFlow**](https://tum-pbs.github.io/PhiFlow/).
+This package is an excellent example of how to interface machine learning and. 
+In fact, it is very close to what we want to have in the future.
+This is a single instantiation of that as their scope is much wider than ours. 
+In addition, we do not offer anything with neural networks but we still try to maintain compatibility.
+
+[**pyqg**](https://github.com/pyqg/pyqg) / [**pyqg-jax**](https://github.com/karlotness/pyqg-jax).
+This package almost works exclusively with the Quasi-Geostrophic equations.
+They also only use the spectral decomposition methodologies.
+We were motivated by community behind this package however we wanted something that was more general that what they offered.
+
+[**jaxdf**](https://github.com/ucl-bug/jaxdf).
+This package was inspirational for defining spatial discretizations.
+
+[**jax-cfd**](https://github.com/google/jax-cfd). 
+This package has been a great inspiration thinking about fluids simulations and machine learning.
+However, we felt that it was very complex and did not cater to people being able to tinker with it at different levels of granularity.
+The internals were optimized but very difficult to real which is a barrier for people with less coding experience.
+
+**Various QG and SW Model Codes**.
+There are many examples of simplified being implemented in Python.
+There are some examples for QG models are [1 Layer QG](https://github.com/redouanelg/qgsw-DI/tree/master), [Spectral QG](https://github.com/hrkz/torchqg), and a [Stacked QG](https://github.com/louity/qgm_pytorch) which are all implemented in PyTorch.
+However, each of them are blueprints.
+Actually, I want even more instances of models being implemented. 
+But we wish for some of them to be integrated into the platform so that we can continue to grow as a community and take the latest and the greatest improvements.
+
+:::
+
+
+
+
+
+
+---
+## Installation
+
+#### pip
+
+We can directly install it via pip from the
+
+```bash
+pip install "git+https://github.com/jejjohnson/jaxsw.git"
+```
+
+#### Cloning
+
+We can also clone the git repository
+
+```bash
+git clone https://github.com/jejjohnson/jaxsw.git
+cd jaxsw
+```
+
+**poetry**
+
+The easiest way to get started is to simply use the poetry package which installs all necessary dev packages as well
 
-## Quick Start
+```bash
+poetry install
+```
 
+**pip**
 
-### **3 APIs for PDEs (TODO)**
+We can also install via `pip` as well
 
-In this 3 part tutorial, we describe three hierarchies for describing PDEs:
+```bash
+pip install .
+```
 
-* **Low-Level**: Fully explicit where we make everything explicit
-* **Mid-Level**: Semi-explicit where we specify some parts but don't specify others.
-* **High-Level**: Fully implicit where we hide many of the mechanics underneath syntactic sugar.
+**Conda**
 
+We also have a conda environment with all of the equivalent dependencies.
 
-### [**Anatomy of a PDE**](./components/anatomy_pde)
+```bash
+conda env create -f environments/jax_linux_cpu.yaml
+conda activate jaxsw
+```
+
+
+---
+
+## Tutorials
+
+
+
+
+#### [**Anatomy of a PDE**](./components/anatomy_pde)
 
 This goes through a full example looking at the components of the `jaxsw` framework.
 We explicitly describe the components of the PDE which will be important for this package.
 This includes:
-1. State
+1. Domain
 2. Params
-3. Geometry
-4. Spatial Discretization
-5. Equation of motion
-6. Time Stepping
+2. Initial Condition Operators
+3. Boundary Operators
+4. Spatial Operators
+5. RHS Operators
+6. Time Steppers
 
 To accomplish this, I will showcase how we can use many other libraries as the backbone to do many canonical PDEs, e.g. [FiniteDiffX](), , [jaxdf](), and [kernex]() for spatial discretizations and [Diffrax]() for timestepping schemes.
 In addition, I will do my best to use some of the best elements of `JAX` to really take advantage of some of the native elements, e.g. `vmap`, `scan`.
 
-### **Spatial Discretization (TODO)**
+---
+
+#### **3 APIs for PDEs (TODO)**
+
+In this 3 part tutorial, we describe three APIs for implementing PDEs:
+
+* **Low-Level/Researcher**: We showcase the functional API which offers a high level of granularity and control
+* **Mid-Level/Engineer**: We showcase the operator API which offers a medium level of granularity and control.
+* **High-Level/Casual**: We showcase the prebuilt models where we need minimum intervention to get started with PDEs.
 
-This tutorial goes through how to do some spatial discretizations using this library.
+---
+
+#### **Spatial Operators (TODO)**
+
+This tutorial goes through how to do some spatial operations using this library.
 We will look at how we can define a simple geometry and then choose various ways to do operations such as finite difference, e.g. slicing, convolutions, or spectral methods.
 In addition, we will look at how we can do some simple procedures to get staggered grids which can greatly improve the accuracy of methods.
 
+---
 
-### **TimeSteppers (TODO)**
+#### [**Grid Operators**](./components/grid_operations.ipynb)
+
+This tutorial goes through how we handle grid operations.
+This is very useful for implementing staggered grids for different fields.
+We showcase how we can use the grid operators to move between the domains along the staggered grid.
+
+---
+
+#### **Boundary Operators (TODO)**
+
+This tutorial goes through how we handle boundary operations.
+These are some of the most important components of PDEs as they essentially govern a huge portion of the dynamics and stability within the system.
+We show case how we can make some custom operators using some of the staple methods like periodic, Dirichlet and Neumann.
+
+---
+
+#### **TimeSteppers (TODO)**
 
 This tutorial goes through how to do time stepping in with JAX.
 I'll show how this can be accomplished from scratch and through the native `JAX` package.
 We also look at the `diffrax` which allows us to remove a lot of the complexity.
 
 
 ---
-## Tutorials
+## Examples
 
-This are more in-depth tutorials about we can use this package for various canonical PDEs that are necessary for understanding and constructing simple differentiable ocean models.
+> This are more in-depth tutorials about we can use this package for various canonical PDEs that are necessary for understanding and constructing simple differentiable ocean models.
 
 #### [**Lorenz ODEs**](./lorenz/lorenz)
 

content/qg/multilayer_qg.ipynb

@@ -8,7 +8,31 @@
     "user_expressions": []
    },
    "source": [
-    "# Multilayer QG"
+    "---\n",
+    "title: 3 Layer Stacked QG Model\n",
+    "subject: Jax Approximate Ocean Models\n",
+    "# subtitle: How can I estimate the state AND the parameters?\n",
+    "short_title: Stacked QG\n",
+    "authors:\n",
+    "  - name: J. Emmanuel Johnson\n",
+    "    affiliations:\n",
+    "      - CNRS\n",
+    "      - MEOM\n",
+    "    orcid: 0000-0002-6739-0053\n",
+    "    email: jemanjohnson34@gmail.com\n",
+    "  - name: Takaya Uchida\n",
+    "    affiliations:\n",
+    "      - FSU\n",
+    "    orcid: https://orcid.org/0000-0002-8654-6009\n",
+    "    email: tuchida@fsu.edu\n",
+    "license: CC-BY-4.0\n",
+    "keywords: jax, shallow water model, differentiable\n",
+    "abbreviations:\n",
+    "    SW: Shallow Water\n",
+    "    QG: Quasi-Geostrophic\n",
+    "    PDE: Partial Differential Equation\n",
+    "    RHS: Right Hand Side\n",
+    "---"
    ]
   },
   {

content/qg/qg_inversion.ipynb

@@ -8,7 +8,31 @@
     "user_expressions": []
    },
    "source": [
-    "# QG Inversion"
+    "---\n",
+    "title: QG Inversion Schemes\n",
+    "subject: Jax Approximate Ocean Models\n",
+    "# subtitle: How can I estimate the state AND the parameters?\n",
+    "short_title: QG Inversion\n",
+    "authors:\n",
+    "  - name: J. Emmanuel Johnson\n",
+    "    affiliations:\n",
+    "      - CNRS\n",
+    "      - MEOM\n",
+    "    orcid: 0000-0002-6739-0053\n",
+    "    email: jemanjohnson34@gmail.com\n",
+    "  - name: Takaya Uchida\n",
+    "    affiliations:\n",
+    "      - FSU\n",
+    "    orcid: https://orcid.org/0000-0002-8654-6009\n",
+    "    email: tuchida@fsu.edu\n",
+    "license: CC-BY-4.0\n",
+    "keywords: jax, shallow water model, differentiable\n",
+    "abbreviations:\n",
+    "    SW: Shallow Water\n",
+    "    QG: Quasi-Geostrophic\n",
+    "    PDE: Partial Differential Equation\n",
+    "    RHS: Right Hand Side\n",
+    "---"
    ]
   },
   {

content/qg/ssh_free_run.ipynb

@@ -9,7 +9,31 @@
     "user_expressions": []
    },
    "source": [
-    "# SSH Free Run"
+    "---\n",
+    "title: Free-Run QG with SSH\n",
+    "subject: Jax Approximate Ocean Models\n",
+    "# subtitle: How can I estimate the state AND the parameters?\n",
+    "short_title: SSH Free-Run QG\n",
+    "authors:\n",
+    "  - name: J. Emmanuel Johnson\n",
+    "    affiliations:\n",
+    "      - CNRS\n",
+    "      - MEOM\n",
+    "    orcid: 0000-0002-6739-0053\n",
+    "    email: jemanjohnson34@gmail.com\n",
+    "  - name: Takaya Uchida\n",
+    "    affiliations:\n",
+    "      - FSU\n",
+    "    orcid: https://orcid.org/0000-0002-8654-6009\n",
+    "    email: tuchida@fsu.edu\n",
+    "license: CC-BY-4.0\n",
+    "keywords: jax, shallow water model, differentiable\n",
+    "abbreviations:\n",
+    "    SW: Shallow Water\n",
+    "    QG: Quasi-Geostrophic\n",
+    "    PDE: Partial Differential Equation\n",
+    "    RHS: Right Hand Side\n",
+    "---"
    ]
   },
   {