diff --git a/docs/.htaccess b/docs/.htaccess new file mode 100644 index 0000000..c8a5446 --- /dev/null +++ b/docs/.htaccess @@ -0,0 +1 @@ +ErrorDocument 404 /equinox/404.html \ No newline at end of file diff --git a/docs/Makefile b/docs/Makefile deleted file mode 100644 index 63a98ea..0000000 --- a/docs/Makefile +++ /dev/null @@ -1,20 +0,0 @@ -# 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) -c . diff --git a/docs/_overrides/partials/source.html b/docs/_overrides/partials/source.html new file mode 100644 index 0000000..10bba66 --- /dev/null +++ b/docs/_overrides/partials/source.html @@ -0,0 +1,20 @@ +{% import "partials/language.html" as lang with context %} + +
+ {% set icon = config.theme.icon.repo or "fontawesome/brands/git-alt" %} + {% include ".icons/" ~ icon ~ ".svg" %} +
+
+ {{ config.repo_name }} +
+ +{% if config.theme.twitter_url %} + +
+ {% include ".icons/fontawesome/brands/twitter.svg" %} +
+
+ {{ config.theme.twitter_name }} +
+ +{% endif %} \ No newline at end of file diff --git a/docs/_static/custom_css.css b/docs/_static/custom_css.css new file mode 100644 index 0000000..e9be9fc --- /dev/null +++ b/docs/_static/custom_css.css @@ -0,0 +1,162 @@ +/* Fix /page#foo going to the top of the viewport and being hidden by the navbar */ +html { + scroll-padding-top: 50px; +} + +/* Fit the Twitter handle alongside the GitHub one in the top right. */ + +div.md-header__source { + width: revert; + max-width: revert; +} + +a.md-source { + display: inline-block; +} + +.md-source__repository { + max-width: 100%; +} + +/* Emphasise sections of nav on left hand side */ + +nav.md-nav { + padding-left: 5px; +} + +nav.md-nav--secondary { + border-left: revert !important; +} + +.md-nav__title { + font-size: 0.9rem; +} + +.md-nav__item--section > .md-nav__link { + font-size: 0.9rem; +} + +/* Indent autogenerated documentation */ + +div.doc-contents { + padding-left: 25px; + border-left: 4px solid rgba(230, 230, 230); +} + +/* Increase visibility of splitters "---" */ + +[data-md-color-scheme="default"] .md-typeset hr { + border-bottom-color: rgb(0, 0, 0); + border-bottom-width: 1pt; +} + +[data-md-color-scheme="slate"] .md-typeset hr { + border-bottom-color: rgb(230, 230, 230); +} + +/* More space at the bottom of the page */ + +.md-main__inner { + margin-bottom: 1.5rem; +} + +/* Remove prev/next footer buttons */ + +.md-footer__inner { + display: none; +} + +/* Change font sizes */ + +html { + /* Decrease font size for overall webpage + Down from 137.5% which is the Material default */ + font-size: 110%; +} + +.md-typeset .admonition { + /* Increase font size in admonitions */ + font-size: 100% !important; +} + +.md-typeset details { + /* Increase font size in details */ + font-size: 100% !important; +} + +.md-typeset h1 { + font-size: 1.6rem; +} + +.md-typeset h2 { + font-size: 1.5rem; +} + +.md-typeset h3 { + font-size: 1.3rem; +} + +.md-typeset h4 { + font-size: 1.1rem; +} + +.md-typeset h5 { + font-size: 0.9rem; +} + +.md-typeset h6 { + font-size: 0.8rem; +} + +/* Bugfix: remove the superfluous parts generated when doing: + +??? Blah + + ::: library.something +*/ + +.md-typeset details .mkdocstrings > h4 { + display: none; +} + +.md-typeset details .mkdocstrings > h5 { + display: none; +} + +/* Change default colours for tags */ + +[data-md-color-scheme="default"] { + --md-typeset-a-color: rgb(0, 189, 164) !important; +} +[data-md-color-scheme="slate"] { + --md-typeset-a-color: rgb(0, 189, 164) !important; +} + +/* Highlight functions, classes etc. type signatures. Really helps to make clear where + one item ends and another begins. */ + +[data-md-color-scheme="default"] { + --doc-heading-color: #DDD; + --doc-heading-border-color: #CCC; + --doc-heading-color-alt: #F0F0F0; +} +[data-md-color-scheme="slate"] { + --doc-heading-color: rgb(25,25,33); + --doc-heading-border-color: rgb(25,25,33); + --doc-heading-color-alt: rgb(33,33,44); + --md-code-bg-color: rgb(38,38,50); +} + +h4.doc-heading { + /* NOT var(--md-code-bg-color) as that's not visually distinct from other code blocks.*/ + background-color: var(--doc-heading-color); + border: solid var(--doc-heading-border-color); + border-width: 1.5pt; + border-radius: 2pt; + padding: 0pt 5pt 2pt 5pt; +} +h5.doc-heading, h6.heading { + background-color: var(--doc-heading-color-alt); + border-radius: 2pt; + padding: 0pt 5pt 2pt 5pt; +} \ No newline at end of file diff --git a/docs/_static/mathjax.js b/docs/_static/mathjax.js new file mode 100644 index 0000000..d883d70 --- /dev/null +++ b/docs/_static/mathjax.js @@ -0,0 +1,16 @@ +window.MathJax = { + tex: { + inlineMath: [["\\(", "\\)"]], + displayMath: [["\\[", "\\]"]], + processEscapes: true, + processEnvironments: true + }, + options: { + ignoreHtmlClass: ".*|", + processHtmlClass: "arithmatex" + } +}; + +document$.subscribe(() => { + MathJax.typesetPromise() +}) \ No newline at end of file diff --git a/docs/api/etdrk/etdrk.md b/docs/api/etdrk/etdrk.md new file mode 100644 index 0000000..dee5fb3 --- /dev/null +++ b/docs/api/etdrk/etdrk.md @@ -0,0 +1,5 @@ +# ETDRK + +::: exponax.etdrk.ETDRK0 + +::: exponax.etdrk.ETDRK1 \ No newline at end of file diff --git a/docs/api/stepper/linear.md b/docs/api/stepper/linear.md new file mode 100644 index 0000000..b04f6f4 --- /dev/null +++ b/docs/api/stepper/linear.md @@ -0,0 +1,3 @@ +# Linear Stepper + +::: exponax.stepper.Advection \ No newline at end of file diff --git a/docs/conf.py b/docs/conf.py deleted file mode 100644 index f23f75e..0000000 --- a/docs/conf.py +++ /dev/null @@ -1,49 +0,0 @@ -# Configuration file for the Sphinx documentation builder. -# -# For the full list of built-in configuration values, see the documentation: -# https://www.sphinx-doc.org/en/master/usage/configuration.html - -# -- Project information ----------------------------------------------------- -# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information - -import re - -from exponax import __version__ # noqa: E402 - -project = "exponax" -copyright = "2024, The exponax Team" -author = "The exponax Team" - -# The short X.Y version -parsed_version = re.match(r"(\d+\.\d+\.\d+)", __version__) -if parsed_version: - version = parsed_version.group(1) -else: - version = "0.0.0" - -# The full version, including alpha/beta/rc tags -release = __version__ - -# -- General configuration --------------------------------------------------- -# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration - -extensions = [ - "sphinx.ext.intersphinx", - "sphinx.ext.autodoc", - "sphinx.ext.napoleon", -] - -intersphinx_mapping = { - "python": ("https://docs.python.org/3", None), - "numpy": ("http://docs.scipy.org/doc/numpy/", None), -} - -templates_path = [] -exclude_patterns = ["build", "Thumbs.db", ".DS_Store"] - - -# -- Options for HTML output ------------------------------------------------- -# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output - -html_theme = "alabaster" -html_static_path = ["static"] diff --git a/docs/index.rst b/docs/index.rst deleted file mode 100644 index 0b2811c..0000000 --- a/docs/index.rst +++ /dev/null @@ -1,6 +0,0 @@ -API documentation -================= - -.. automodule:: exponax - :members: - :undoc-members: \ No newline at end of file diff --git a/docs/make.bat b/docs/make.bat deleted file mode 100644 index 32bb245..0000000 --- a/docs/make.bat +++ /dev/null @@ -1,35 +0,0 @@ -@ECHO OFF - -pushd %~dp0 - -REM Command file for Sphinx documentation - -if "%SPHINXBUILD%" == "" ( - set SPHINXBUILD=sphinx-build -) -set SOURCEDIR=. -set BUILDDIR=_build - -%SPHINXBUILD% >NUL 2>NUL -if errorlevel 9009 ( - echo. - echo.The 'sphinx-build' command was not found. Make sure you have Sphinx - echo.installed, then set the SPHINXBUILD environment variable to point - echo.to the full path of the 'sphinx-build' executable. Alternatively you - echo.may add the Sphinx directory to PATH. - echo. - echo.If you don't have Sphinx installed, grab it from - echo.https://www.sphinx-doc.org/ - exit /b 1 -) - -if "%1" == "" goto help - -%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% -goto end - -:help -%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% - -:end -popd diff --git a/docs/requirements.txt b/docs/requirements.txt index 45a6471..fba2309 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -1 +1,14 @@ -sphinx==7.0.1 \ No newline at end of file +mkdocs==1.3.0 # Main documentation generator. +mkdocs-material==7.3.6 # Theme +pymdown-extensions==9.4 # Markdown extensions e.g. to handle LaTeX. +mkdocstrings==0.17.0 # Autogenerate documentation from docstrings. +mknotebooks==0.7.1 # Turn Jupyter Lab notebooks into webpages. +pytkdocs_tweaks==0.0.8 # Tweaks mkdocstrings to improve various aspects +mkdocs_include_exclude_files==0.0.1 # Tweak which files are included/excluded +jinja2==3.0.3 # Older version. After 3.1.0 seems to be incompatible with current versions of mkdocstrings. +pygments==2.14.0 +nbconvert==6.5.0 # Fix https://github.com/patrick-kidger/equinox/issues/428 + +jax[cpu] +equinox +matplotlib \ No newline at end of file diff --git a/docs/static/.empty b/docs/static/.empty deleted file mode 100644 index e69de29..0000000 diff --git a/mkdocs.yml b/mkdocs.yml new file mode 100644 index 0000000..c89c75b --- /dev/null +++ b/mkdocs.yml @@ -0,0 +1,153 @@ +theme: + name: material + features: + - navigation.sections # Sections are included in the navigation on the left. + - toc.integrate # Table of contents is integrated on the left; does not appear separately on the right. + - header.autohide # header disappears as you scroll + palette: + # Light mode / dark mode + # We deliberately don't automatically use `media` to check a user's preferences. We default to light mode as + # (a) it looks more professional, and (b) is more obvious about the fact that it offers a (dark mode) toggle. + - scheme: default + primary: white + accent: amber + toggle: + icon: material/weather-night + name: Switch to dark mode + - scheme: slate + primary: black + accent: amber + toggle: + icon: material/weather-sunny + name: Switch to light mode + icon: + repo: fontawesome/brands/github # GitHub logo in top right + logo: "material/circle-opacity" # Equinox logo in top left + favicon: "_static/favicon.png" + custom_dir: "docs/_overrides" # Overriding part of the HTML + + # These additions are my own custom ones, having overridden a partial. + twitter_name: "@felix_m_koehler" + twitter_url: "https://twitter.com/felix_m_koehler" + +site_name: Exponax +site_description: The documentation for the Exponax software. +site_author: Felix Koehler +site_url: https://fkoehler.site/exponax + +repo_url: https://github.com/ceyron/exponax +repo_name: ceyron/exponax +edit_uri: "" # No edit button, as some of our pages are in /docs and some in /examples via symlink, so it's impossible for them all to be accurate + +strict: true # Don't allow warnings during the build process + +extra_javascript: + # The below three make MathJax work, see https://squidfunk.github.io/mkdocs-material/reference/mathjax/ + - _static/mathjax.js + - https://polyfill.io/v3/polyfill.min.js?features=es6 + - https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js + +extra_css: + - _static/custom_css.css + +markdown_extensions: + - pymdownx.arithmatex: # Render LaTeX via MathJax + generic: true + - pymdownx.superfences # Seems to enable syntax highlighting when used with the Material theme. + - pymdownx.details # Allowing hidden expandable regions denoted by ??? + - pymdownx.snippets: # Include one Markdown file into another + base_path: docs + - admonition + - toc: + permalink: "ยค" # Adds a clickable permalink to each section heading + toc_depth: 4 + +plugins: + - search # default search plugin; needs manually re-enabling when using any other plugins + - autorefs # Cross-links to headings + - include_exclude_files: + include: + - ".htaccess" + exclude: + - "_overrides" + - "examples/MNIST" + - "examples/bert_checkpoint.eqx" + - mknotebooks # Jupyter notebooks + - mkdocstrings: + handlers: + python: + setup_commands: + - import pytkdocs_tweaks + - pytkdocs_tweaks.main() + - import jaxtyping + - jaxtyping.set_array_name_format("array") + - import jax + - jax.ShapeDtypeStruct.__module__ = "jax" + - jax.core.ClosedJaxpr.__module__ = "jax.core" + + selection: + inherited_members: true # Allow looking up inherited methods + rendering: + show_root_heading: true # actually display anything at all... + show_root_full_path: true # display "diffrax.asdf" not just "asdf" + show_if_no_docstring: true + show_signature_annotations: true + show_source: false # don't include source code + members_order: source # order methods according to their order of definition in the source code, not alphabetical order + heading_level: 4 + +nav: + # - 'index.md' + # - 'all-of-equinox.md' + # - Examples: + # - Introductory: + # - CNN on MNIST: 'examples/mnist.ipynb' + # - Train RNN: 'examples/train_rnn.ipynb' + # - Advanced: + # - Generative score-based diffusion: 'examples/score_based_diffusion.ipynb' + # - BERT language model: 'examples/bert.ipynb' + # - U-Net implementation: 'examples/unet.ipynb' + # - Vision transformer: 'examples/vision_transformer.ipynb' + # - Image GAN: 'examples/deep_convolutional_gan.ipynb' + # - Mistral 7B: 'examples/mistral_7b.md' + # - Features: + # - Freezing parameters: 'examples/frozen_layer.ipynb' + # - Compatibility with init-apply libraries: 'examples/init_apply.ipynb' + # - Stateful operations (e.g. BatchNorm): 'examples/stateful.ipynb' + # - Autoparallelism (e.g. multi-GPU): 'examples/parallelism.ipynb' + # - Serialisation (with hyperparameters): 'examples/serialisation.ipynb' + - Basic API: + - ETDRK: + - 'api/etdrk/etdrk.md' + - Stepper: + - 'api/stepper/linear.md' + # - 'api/module/advanced_fields.md' + # - Neural network layers: + # - 'api/nn/linear.md' + # - 'api/nn/conv.md' + # - 'api/nn/rnn.md' + # - 'api/nn/attention.md' + # - 'api/nn/activations.md' + # - 'api/nn/pool.md' + # - 'api/nn/dropout.md' + # - 'api/nn/normalisation.md' + # - 'api/nn/embedding.md' + # - 'api/nn/mlp.md' + # - 'api/nn/sequential.md' + # - 'api/nn/inference.md' + # - 'api/nn/shared.md' + # - 'api/nn/stateful.md' + # - 'api/transformations.md' + # - 'api/manipulation.md' + # - Advanced API: + # - 'api/caches.md' + # - 'api/debug.md' + # - 'api/enumerations.md' + # - 'api/errors.md' + # - 'api/pretty-printing.md' + # - 'api/serialisation.md' + # - Misc: + # - 'faq.md' + # - 'tricks.md' + # # - 'pattern.md' + # - 'citation.md' \ No newline at end of file