Scientific Coding

Prompt

**System Instructions / Persona:** Act as an expert in JAX, scientific computing, and computational astrophysics. **Objective:** Write a complete, high-performance **JAX-based adaptive mesh N-body simulation**. **Technical & Architectural Requirements:** 1. **Differentiability:** The entire simulation step must be fully end-to-end differentiable using JAX (`jax.grad`, `jax.vmap`, `jax.jit`). 2. **Custom PyTree Data Structure:** * Define a unified custom PyTree node (using `jax.tree_util.register_dataclass`) that encapsulates both the **adaptive mesh grid** and the **particle state** (positions, velocities, masses). * Ensure tree flatten/unflatten methods are correctly implemented if using standard JAX syntax. 3. **Adaptive Mesh Logic:** * Implement functional/differentiable logic for mesh refinement based on local particle density. * Include standard particle-mesh steps: Particle-to-Mesh mass deposition (e.g., Cloud-in-Cell or Nearest-Grid-Point), solving Poisson's equation or mesh-based force calculations, and Mesh-to-Particle force interpolation. 4. **Primary Entry Function:** * Provide a clean, unified top-level function with the exact signature: ```python def simulate(state: SimulationState, num_steps: int, dt: float, **kwargs) -> SimulationState: ``` * Use `jax.lax.scan` internally for time-stepping to maintain compatibility with `jax.jit` and forward/reverse-mode automatic differentiation. **Output Deliverables:** * **Self-Contained Python Code:** Fully executable with all necessary imports (`jax`, `jax.numpy`, etc.). * **Example Execution & Differentiation Verification:** At the end of the response, include a functional example run demonstrating: 1. Initialization of particles and mesh structure. 2. Running `simulate(...)`. 3. Computing the gradient of a scalar metric (e.g., final kinetic energy or center-of-mass dispersion) with respect to initial particle conditions using `jax.grad(simulate)`. --- ### Key Improvements Made to Your Prompt: * **Explicit JAX Requirements (`jax.lax.scan`):** Standard Python loops break or slow down execution inside JAX jit/grad routines. Specifying `lax.scan` forces the model to write clean, traceable code. * **Defined Data Structures:** Specifying how to register custom PyTrees ensures the model doesn't use unsupported Python dicts or standard classes that cause JAX tracer errors. * **Functional Adaptive Mesh Steps:** Clearly lists the required physics components (Particle-to-Mesh, solver, Mesh-to-Particle) so the model doesn't just produce placeholder dummy code. * **Verification via `jax.grad`:** Requesting a working gradient check at the end forces the model to ensure its mesh implementation is actually differentiable, rather than using non-differentiable operations like in-place array updates or raw Python branches.