Skip to content

feat: add code for drawing full scenes#209

Draft
beckermr wants to merge 8 commits intomainfrom
fix-bounds-scene-rend
Draft

feat: add code for drawing full scenes#209
beckermr wants to merge 8 commits intomainfrom
fix-bounds-scene-rend

Conversation

@beckermr
Copy link
Collaborator

@beckermr beckermr commented Mar 11, 2026

This PR addresses some of the full-scene rendering issues raised in #190.

  • bounds comparisons can be vmapped
  • can render objects in stamps and then apply to full image via vmapping
  • can render sum of objects into the full image

There were a few issues, many of which were unrelated:

  • some bugs in the repr functions have been fixed
  • the bounds objects' pytrees change shape when traced
  • images need to tell JAX to trace their attached bounds objects (these were marked as static data)
  • We cannot render objects with the center keyword set to a JAX value since it uses the bounds to basically put the object at the right location in the stamp. So instead, I coded up rendering with an offset and applying the bounds shifts after.
  • For JIT-ed code, we need dynamic slice operators to assemble the final image.
  • The final bit is that for JAX you need fixed size arrays. So when you assemble the stamps, you actually need to pad the image you are adding into by the stamp size and then remove the padding. This avoids having to add a dynamically sized array into the final image.

One change that we could explore to make all of this more automatic would be to represent image bounds as a location + width, instead of (xmin, xmax, ymin, ymax). This would be more compatible with JAX since we could force the width to be a python int and then automatically use dynamic slice operators to add images together within the bounds.

@codspeed-hq
Copy link

codspeed-hq bot commented Mar 11, 2026

Merging this PR will degrade performance by 48.3%

⚠️ 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.

⚠️ Different runtime environments detected

Some benchmarks with significant performance changes were compared across different runtime environments,
which may affect the accuracy of the results.

Open the report in CodSpeed to investigate

⚡ 3 improved benchmarks
❌ 4 regressed benchmarks
✅ 29 untouched benchmarks

⚠️ Please fix the performance issues or acknowledge them on CodSpeed.

Performance Changes

Mode Benchmark BASE HEAD Efficiency
WallTime test_benchmark_spergel_conv[run] 265.2 ms 382.6 ms -30.68%
WallTime test_benchmarks_lanczos_interp[xval-conserve_dc-run] 114.9 µs 85.3 µs +34.62%
WallTime test_benchmarks_metacal[run] 18.4 ms 35.6 ms -48.3%
WallTime test_benchmarks_lanczos_interp[xval-no_conserve_dc-run] 101.4 µs 76 µs +33.42%
WallTime test_benchmark_moffat_init[run] 83.1 µs 62.3 µs +33.42%
WallTime test_benchmark_spergel_xvalue[run] 1.6 s 2.9 s -42.61%
WallTime test_benchmark_moffat_conv[run] 290.1 ms 423.9 ms -31.58%

Comparing fix-bounds-scene-rend (a401286) with main (d2f2d94)

Open in CodSpeed

@beckermr beckermr changed the title fix: ensure bounds comparisons can be vmapped feat: add code for drawing full scenes Mar 13, 2026
return int(x)
except Exception:
try:
if not jnp.any(jnp.isnan(x)):
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All the changes in this block need to be backed out and redone in a new way. This was hack to get the thing moving. The issue is that sometimes this function encounters not just arrays, floats, ints, and tracers, but also jax leaf objects. :/

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant