WIP: Add documentation site with Material for MkDocs by EiffL · Pull Request #204 · GalSim-developers/JAX-GalSim

EiffL · 2026-02-20T17:59:33Z

Summary

Adds a full documentation site using Material for MkDocs with mkdocstrings for auto-generated API reference from docstrings
Includes getting started guides (installation, quick start), notable differences from GalSim, API coverage page, and complete API reference stubs for all implemented modules
Adds GitHub Actions workflow for building docs on PRs and deploying to GitHub Pages on push to main

What's included

Home page with feature highlights, quick install, and minimal code example
Getting Started section: installation (pip, GPU, dev setup) and quick start tutorial with jit/grad/vmap examples
Notable Differences from GalSim: thorough guide covering immutability, array views, RNG, PyTree registration, control flow/tracing, profile restrictions, numerical precision, and the @implements decorator
API Coverage page (22.5% of GalSim API)
API Reference (~30 pages) auto-generated from docstrings via mkdocstrings
CI workflow (.github/workflows/docs.yml): mkdocs build --strict on PRs, mkdocs gh-deploy on push to main

Known limitations

API function documentation is incomplete — the @implements decorator copies GalSim's docstrings at runtime, but mkdocstrings doesn't always pick them up cleanly. RST cross-references (:func:, :class:) from GalSim's docstrings render as literal text rather than links.
Some API pages show minimal content where docstring propagation from GalSim didn't work as expected.

🤖 Generated with Claude Code

Set up full documentation infrastructure with auto-generated API reference from docstrings, narrative guides, and GitHub Pages deployment via CI. - Add mkdocs-material and mkdocstrings[python] as docs dependencies - Create mkdocs.yml with Material theme, left sidebar nav, Mermaid, MathJax - Add narrative pages: home, getting started, architecture, notable differences - Add API reference stubs for all public modules targeting specific classes - Add GitHub Actions workflow for build on PR / deploy on push to main Co-Authored-By: Claude Opus 4.6 <[email protected]>

Remove duplicated explanations, tighten prose, and condense sections that repeated information already covered on other pages. Net reduction of ~80 lines with no content loss -- cross-links replace duplication. Co-Authored-By: Claude Opus 4.6 <[email protected]>

@implements

Remove the architecture/ pages (pytree, implements, gsobject, drawing) and key-concepts page. Consolidate their most useful content into notable-differences.md, which now thoroughly covers immutability, array views, RNG, PyTree registration, control flow/tracing, profile restrictions, numerical precision, and the @implements decorator. Co-Authored-By: Claude Opus 4.6 <[email protected]>

codspeed-hq · 2026-02-20T18:13:34Z

Merging this PR will not alter performance

⚠️

Unknown Walltime execution environment detected

Using the Walltime instrument on standard Hosted Runners will lead to inconsistent data.

For the most accurate results, we recommend using CodSpeed Macro Runners: bare-metal machines fine-tuned for performance measurement consistency.

✅ 36 untouched benchmarks

_{Comparing documentation (5b2c943) with main (d2f2d94)}

beckermr · 2026-02-20T19:59:33Z

docs/api-coverage.md

+
+## Supported APIs
+
+??? note "Click to expand the full list of implemented APIs"


Is this AI slop or actual markdown.

I think it's actually syntax from the admonition plugin? https://squidfunk.github.io/mkdocs-material/reference/admonitions/#collapsible-blocks

ismael-mendoza

I flagged a couple of cases in the examples provided that don't actually work. @EiffL could you modify as needed?

ismael-mendoza · 2026-03-04T01:11:37Z

docs/notable-differences.md

+image.array[10, 10] = 0.0
+
+# JAX-GalSim — returns a new image each time
+image = image.addNoise(noise)


is this actually the case? looking at the addNoise methods in noise.py, I don't think they return the image object.

ismael-mendoza · 2026-03-04T01:19:07Z

docs/notable-differences.md

+sequences differ from GalSim. Results will not match GalSim number-for-number.
+This is expected --- the underlying PRNG algorithms are completely different.
+
+**No in-place fill**: GalSim deviates can "fill" existing arrays. JAX deviates


This feels like it should be in a different section, also again like above not sure if it's true in all cases.

ismael-mendoza · 2026-03-04T01:36:01Z

docs/getting-started/quickstart.md

+```python
+import jax
+
+@jax.jit
+def simulate(flux, sigma):
+    gal = jax_galsim.Gaussian(flux=flux, sigma=sigma)
+    psf = jax_galsim.Gaussian(flux=1.0, sigma=1.0)
+    final = jax_galsim.Convolve([gal, psf])
+    return final.drawImage(scale=0.2)
+
+# First call compiles; subsequent calls are fast
+image = simulate(1e5, 2.0)
+```


this does not work without setting the gsparams with fft_size and specifying the size of the image to draw

ismael-mendoza · 2026-03-04T01:39:04Z

docs/getting-started/quickstart.md

+```python
+import jax.numpy as jnp
+
+sigmas = jnp.linspace(1.0, 4.0, 10)
+
+@jax.vmap
+def batch_simulate(sigma):
+    gal = jax_galsim.Gaussian(flux=1e5, sigma=sigma)
+    psf = jax_galsim.Gaussian(flux=1.0, sigma=1.0)
+    final = jax_galsim.Convolve([gal, psf])
+    return final.drawImage(scale=0.2, nx=64, ny=64).array
+
+# Simulate all 10 galaxies in parallel
+images = batch_simulate(sigmas)  # shape: (10, 64, 64)
+```


this example also requires specifying the gsparams

ismael-mendoza · 2026-03-04T01:49:26Z

docs/index.md

+def simulate(flux, sigma):
+    gal = jax_galsim.Gaussian(flux=flux, sigma=sigma)
+    psf = jax_galsim.Gaussian(flux=1.0, sigma=1.0)
+    return jax_galsim.Convolve([gal, psf]).drawImage(scale=0.2).array.sum()


need to specify fft_size and image size explicitly for jiiting to work

ismael-mendoza · 2026-03-04T01:51:23Z

docs/notable-differences.md

+```python
+import jax
+
+gsparams = jax_galsim.GSParams(maximum_fft_size=8192)
+
+@jax.jit
+def simulate(flux, sigma):
+    gal = jax_galsim.Gaussian(flux=flux, sigma=sigma, gsparams=gsparams)
+    return gal.drawImage(scale=0.2).array.sum()


this example will not work either, see other comments

ismael-mendoza · 2026-03-04T01:52:24Z

docs/notable-differences.md

+@jax.vmap
+def batch(sigma):
+    gal = jax_galsim.Gaussian(flux=1e5, sigma=sigma)
+    # Must specify nx, ny so all images have the same shape
+    return gal.drawImage(scale=0.2, nx=64, ny=64).array


this example also won't work, need gsparams

ismael-mendoza · 2026-03-04T01:59:41Z

docs/notable-differences.md

+
+```python
+noise = jax_galsim.GaussianNoise(sigma=30.0)
+image = image.addNoise(noise)  # state is managed internally


addNoise does not actually return a new image object

ismael-mendoza · 2026-03-04T02:00:58Z

docs/index.md

+image = final.drawImage(scale=0.2)
+
+# Add noise
+image = image.addNoise(jax_galsim.GaussianNoise(sigma=30.0))


add noise does not return a new object, it modifies it in place like the OG

EiffL and others added 3 commits February 20, 2026 17:52

EiffL changed the title ~~Add documentation site with Material for MkDocs~~ WIP: Add documentation site with Material for MkDocs Feb 20, 2026

beckermr reviewed Feb 20, 2026

View reviewed changes

ismael-mendoza requested changes Mar 4, 2026

View reviewed changes

Merge branch 'main' into documentation

5b2c943


		## Supported APIs

		??? note "Click to expand the full list of implemented APIs"

Conversation

EiffL commented Feb 20, 2026

Summary

What's included

Known limitations

Uh oh!

codspeed-hq bot commented Feb 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merging this PR will not alter performance

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ismael-mendoza left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

codspeed-hq bot commented Feb 20, 2026 •

edited

Loading