Docs improvements

.github/CONTRIBUTING.md

@@ -33,10 +33,9 @@ pip install -e .[dev]
 * Docstrings: [pep257](https://www.python.org/dev/peps/pep-0257/) & [NumPy style](https://numpydoc.readthedocs.io/en/latest/format.html).
 * Type hints: [module documentation](https://docs.python.org/3/library/typing.html), [cheat sheet](https://mypy.readthedocs.io/en/stable/cheat_sheet_py3.html).
 
-## Autograd
-This project uses [`autograd`](https://github.com/HIPS/autograd) for automatic differentiation.
-Autograd does not support all NumPy and SciPy functionalities, see [autograd documentation](https://github.com/HIPS/autograd/blob/master/docs/tutorial.md#supported-and-unsupported-parts-of-numpyscipy).
-*NOTE:* using unsupported functionalities results in the gradient calculation failing silently.
+## Jax
+This project uses [`Jax`](https://github.com/jax-ml/jax) for automatic differentiation.
+Jax does not support all NumPy and SciPy functionalities, see [jax documentation](https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html).
 
 ## Pull Requests
   1. Create a fork of WecOptTool

.gitignore

@@ -17,3 +17,5 @@ results*
 */.ipynb_checkpoints/*
 docs/source/_static/theory_animation_ps.gif
 docs/source/_static/theory_animation_td.gif
+docs/source/_static/pseudospectral_equivalence.png
+docs/source/_static/wavebot_ex_xopt_iterates.gif

docs/build_docs.py

@@ -233,6 +233,7 @@ def main() -> None:
         logger.info("Building docs for current branch in debug mode")
         build_doc("latest", build)
         move_pages_debug()
+        print(f"\nThe final HTML pages are in {os.path.join(docs_dir, 'pages')}.")
         return
 
     versions = _load_versions()

docs/clean_docs.py

@@ -22,7 +22,10 @@ def clean() -> None:
         _assert_within_docs(target)
         if os.path.exists(target):
             print(f"Removing {target}")
-            rmtree(target, ignore_errors=False)
+            try:
+                rmtree(target, ignore_errors=False)
+            except Exception as err:
+                print(f"Failed to remove {target}: {err}")
         else:
             print(f"Skipping {target} (not found)")
 

docs/source/conf.py

@@ -35,6 +35,27 @@
     versions = yaml.safe_load(v_file)
 
 
+# pandoc setup
+def _ensure_pandoc_available() -> None:
+    if shutil.which("pandoc") is not None:
+        return
+
+    try:
+        import pypandoc
+        pandoc_path = pypandoc.get_pandoc_path()
+    except Exception as exc:
+        raise RuntimeError(
+            "Pandoc executable was not found. Install it with "
+            "`conda install -c conda-forge pandoc` or install the dev "
+            "dependencies with `pypandoc-binary` included."
+        ) from exc
+
+    pandoc_dir = os.path.dirname(pandoc_path)
+    os.environ["PATH"] = pandoc_dir + os.pathsep + os.environ.get("PATH", "")
+
+_ensure_pandoc_available()
+
+
 def _normalize_ref(ref: str | None) -> str | None:
     if ref is None:
         return None
@@ -94,6 +115,10 @@ def _resolve_version_context(config=None, current_ref_override=None) -> tuple[st
 
 
 skip_notebook_execution = os.environ.get('WOT_DOCS_SKIP_NOTEBOOK_EXECUTION') == '1'
+skip_theory_animations = os.environ.get('WOT_DOCS_SKIP_THEORY_ANIMATIONS') == '1'
+skip_pseudospectral_visualizations = os.environ.get(
+    'WOT_DOCS_SKIP_PSEUDOSPECTRAL_VISUALIZATIONS'
+) == '1'
 
 # -- General configuration ---------------------------------------------------
 extensions = [
@@ -137,17 +162,19 @@ def _all_but_nc(_dir, contents):
 def _copy_examples() -> None:
     print('Copy example notebooks into docs/_examples')
     examples_dst = os.path.join(project_root, 'docs/source/_examples')
-    shutil.rmtree(examples_dst, ignore_errors=True)
+    os.makedirs(examples_dst, exist_ok=True)
+
     shutil.copytree(
         os.path.join(project_root, 'examples'),
         examples_dst,
         ignore=_all_but_ipynb,
-        dirs_exist_ok=True
+        dirs_exist_ok=True,
     )
     shutil.copytree(
         os.path.join(project_root, 'examples/data'),
-        os.path.join(project_root, 'docs/source/_examples/data'),
+        os.path.join(examples_dst, 'data'),
         ignore=_all_but_nc,
+        dirs_exist_ok=True,
     )
 
 
@@ -165,6 +192,20 @@ def _generate_theory_animations() -> None:
     _theory_animations_generated = True
 
 
+def _generate_pseudospectral_visualizations() -> None:
+    global _pseudospectral_visualizations_generated
+    if _pseudospectral_visualizations_generated:
+        return
+
+    importlib.invalidate_caches()
+    module_name = 'make_pseudospectral_visualizations'
+    if module_name in sys.modules:
+        importlib.reload(sys.modules[module_name])
+    else:
+        importlib.import_module(module_name)
+    _pseudospectral_visualizations_generated = True
+
+
 def _cleanup_index_html(outdir: str) -> None:
     index_file = os.path.join(outdir, 'index.html')
     if not os.path.exists(index_file):
@@ -195,8 +236,14 @@ def _on_config_inited(_app, _config):
 
     if skip_notebook_execution:
         print('Skipping notebook execution')
-
-    _generate_theory_animations()
+    if skip_theory_animations:
+        print('Skipping theory animation generation')
+    else:
+        _generate_theory_animations()
+    if skip_pseudospectral_visualizations:
+        print('Skipping pseudospectral visualization generation')
+    else:
+        _generate_pseudospectral_visualizations()
 
 
 def _on_build_finished(app, exception):
@@ -212,6 +259,7 @@ def setup(app):
 
 
 _theory_animations_generated = False
+_pseudospectral_visualizations_generated = False
 
 suppress_warnings = ['autosectionlabel.*', # nbsphinx and austosectionlabel do not play well together
                      'app.add_node', # using multiple builders in custom Sphinx objects throws a bunch of these