Code update

PiperOrigin-RevId: 618464158

docs/Makefile

@@ -0,0 +1,19 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line.
+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/api.rst

@@ -0,0 +1,14 @@
+######################
+Swirl-LM API reference
+######################
+
+.. currentmodule:: swirl_lm.base.driver
+
+.. autofunction:: solver
+
+.. currentmodule:: swirl_lm.base.parameters
+
+.. autofunction::  params_from_config_file_flag
+
+.. autoclass:: SwirlLMParameters
+   :members:

docs/cloud_tpu.md

@@ -43,7 +43,7 @@ machine running the browser to Jupyter.
     TPU_FLAGS=--xla_jf_conv_full_precision=true
 
     gcloud compute tpus tpu-vm create $TPU --project=$PROJECT --zone=$ZONE \
-    --accelerator-type=$TYPE --version=tpu-vm-tf-2.12.0-pod \
+    --accelerator-type=$TYPE --version=tpu-vm-tf-2.14.0-pod \
     --metadata="tensorflow-env-vars=-e LIBTPU_INIT_ARGS=$TPU_FLAGS"
     ```
 
@@ -162,7 +162,7 @@ pod. For example, to set xla_jf_conv_full_precision to true:
 
 ```shell
 gcloud compute tpus tpu-vm create $TPU --project=$PROJECT --zone=$ZONE \
-  --accelerator-type=v3-32 --version=tpu-vm-tf-2.12.0-pod \
+  --accelerator-type=v3-32 --version=tpu-vm-tf-2.14.0-pod \
   --metadata="tensorflow-env-vars=-e LIBTPU_INIT_ARGS=--xla_jf_conv_full_precision=true"
 ```
 

docs/conf.py

@@ -0,0 +1,153 @@
+# Copyright 2024 The swirl_lm Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""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:
+# http://www.sphinx-doc.org/en/master/config
+
+# -- 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.
+
+# pylint: disable=g-bad-import-order
+# pylint: disable=g-import-not-at-top
+import doctest
+import inspect
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath('..'))
+sys.path.append(os.path.abspath('ext'))
+
+import swirl_lm
+
+from sphinxcontrib import katex
+
+# -- Project information -----------------------------------------------------
+
+project = 'Swirl-LM'
+copyright = '2024, Google Research'  # pylint: disable=redefined-builtin
+author = 'Swirl-LM Contributors'
+
+# -- General configuration ---------------------------------------------------
+
+master_doc = 'index'
+
+# 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.autosummary',
+    'sphinx.ext.doctest',
+    'sphinx.ext.inheritance_diagram',
+    'sphinx.ext.linkcode',
+    'sphinx.ext.napoleon',
+    'sphinxcontrib.katex',
+    'sphinx_autodoc_typehints',
+]
+
+# 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 autodoc -----------------------------------------------------
+
+autodoc_default_options = {
+    'member-order': 'bysource',
+    'special-members': True,
+    'exclude-members': '__repr__, __str__, __weakref__',
+}
+
+# -- 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 = []
+html_favicon = None
+
+# -- Options for doctest -----------------------------------------------------
+
+doctest_test_doctest_blocks = 'true'
+doctest_global_setup = ''
+doctest_default_flags = (
+    doctest.ELLIPSIS
+    | doctest.IGNORE_EXCEPTION_DETAIL
+    | doctest.DONT_ACCEPT_TRUE_FOR_1
+    | doctest.NORMALIZE_WHITESPACE)
+
+# -- Options for katex ------------------------------------------------------
+
+# See: https://sphinxcontrib-katex.readthedocs.io/en/0.4.1/macros.html
+latex_macros = r"""
+    \def \d              #1{\operatorname{#1}}
+"""
+
+# Translate LaTeX macros to KaTeX and add to options for HTML builder
+katex_macros = katex.latex_defs_to_katex_macros(latex_macros)
+katex_options = 'macros: {' + katex_macros + '}'
+
+# Add LaTeX macros for LATEX builder
+latex_elements = {'preamble': latex_macros}
+
+# -- Source code links -------------------------------------------------------
+
+
+def linkcode_resolve(domain, info):
+  """Resolve a GitHub URL corresponding to Python object."""
+  if domain != 'py':
+    return None
+
+  try:
+    mod = sys.modules[info['module']]
+  except ImportError:
+    return None
+
+  obj = mod
+  try:
+    for attr in info['fullname'].split('.'):
+      obj = getattr(obj, attr)
+  except AttributeError:
+    return None
+  else:
+    obj = inspect.unwrap(obj)
+
+  try:
+    filename = inspect.getsourcefile(obj)
+  except TypeError:
+    return None
+
+  try:
+    source, lineno = inspect.getsourcelines(obj)
+  except OSError:
+    return None
+
+  return ('https://github.com/google-research/swirl-lm/blob/main/swirl_lm/%s'
+          '#L%d#L%d' % (
+              os.path.relpath(filename, start=os.path.dirname(
+                  swirl_lm.__file__)), lineno, lineno + len(source) - 1))

docs/index.rst

@@ -0,0 +1,15 @@
+:github_url: https://github.com/google-research/swirl-lm
+
+Swirl-LM Documentation
+=====================
+
+.. toctree::
+   :caption: API reference
+   :maxdepth: 2
+
+   api
+
+Indices and tables
+==================
+
+* :ref:`genindex`

swirl_lm/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The swirl_lm Authors.
+# Copyright 2024 The swirl_lm Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

swirl_lm/base/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The swirl_lm Authors.
+# Copyright 2024 The swirl_lm Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.