|
| 1 | +# Copyright 2025 Pasteur Labs. All Rights Reserved. |
| 2 | +# SPDX-License-Identifier: Apache-2.0 |
| 3 | + |
| 4 | +"""Single-timestep Burgers' equation solver Tesseract (PyTorch). |
| 5 | +
|
| 6 | +Solves one explicit Euler step of the 1D viscous Burgers' equation: |
| 7 | +
|
| 8 | + u^{n+1} = u^n + dt * (-u * du/dx + nu * d²u/dx²) |
| 9 | +
|
| 10 | +The viscosity field nu is provided as an input — the solver does not compute it. |
| 11 | +This clean interface (state + material field → next state) is the same contract |
| 12 | +that a Fortran solver with an adjoint could implement. The outer time-stepping |
| 13 | +loop and closure evaluation live in the caller, enabling per-timestep closure |
| 14 | +calls and end-to-end gradient flow through both solver and closure. |
| 15 | +""" |
| 16 | + |
| 17 | +from typing import Any |
| 18 | + |
| 19 | +import numpy as np |
| 20 | +import torch |
| 21 | +from pydantic import BaseModel, Field |
| 22 | +from torch.utils._pytree import tree_map |
| 23 | + |
| 24 | +from tesseract_core.runtime import Array, Differentiable, Float64 |
| 25 | +from tesseract_core.runtime.tree_transforms import filter_func, flatten_with_paths |
| 26 | + |
| 27 | +# Default grid size |
| 28 | +N = 128 |
| 29 | + |
| 30 | +# --- Grid setup (fixed for this Tesseract) --- |
| 31 | +DX = 1.0 / (N - 1) |
| 32 | + |
| 33 | +to_tensor = lambda x: ( |
| 34 | + torch.tensor(x, dtype=torch.float64) |
| 35 | + if isinstance(x, np.generic | np.ndarray) |
| 36 | + else x |
| 37 | +) |
| 38 | + |
| 39 | + |
| 40 | +class InputSchema(BaseModel): |
| 41 | + u: Differentiable[Array[(N,), Float64]] = Field( |
| 42 | + description="Current velocity field on the grid" |
| 43 | + ) |
| 44 | + nu: Differentiable[Array[(N,), Float64]] = Field( |
| 45 | + description="Viscosity field at each grid point (must be positive)" |
| 46 | + ) |
| 47 | + dt: float = Field(description="Time step size", default=1e-4) |
| 48 | + |
| 49 | + |
| 50 | +class OutputSchema(BaseModel): |
| 51 | + u_next: Differentiable[Array[(N,), Float64]] = Field( |
| 52 | + description="Velocity field after one time step" |
| 53 | + ) |
| 54 | + |
| 55 | + |
| 56 | +def evaluate(inputs: dict) -> dict: |
| 57 | + """Core differentiable computation — pure torch operations.""" |
| 58 | + u = inputs["u"] |
| 59 | + nu = inputs["nu"] |
| 60 | + dt = inputs["dt"] |
| 61 | + |
| 62 | + # Spatial derivatives via plain central differences. The flow stays smooth |
| 63 | + # and low-Reynolds here (gentle low-frequency ICs, viscosity bounded by the |
| 64 | + # closure's sigmoid so nu*dt/dx^2 stays within the explicit CFL limit), so no |
| 65 | + # shocks form and we don't need upwinding, a conservative/flux form, or a |
| 66 | + # stiff integrator (e.g. ETDRK). A real solver in a harder regime would. |
| 67 | + dudx = torch.zeros_like(u) |
| 68 | + dudx[1:-1] = (u[2:] - u[:-2]) / (2 * DX) |
| 69 | + |
| 70 | + d2udx2 = torch.zeros_like(u) |
| 71 | + d2udx2[1:-1] = (u[2:] - 2 * u[1:-1] + u[:-2]) / (DX**2) |
| 72 | + |
| 73 | + # Burgers' equation: du/dt = -u * du/dx + nu * d²u/dx² |
| 74 | + dudt = -u * dudx + nu * d2udx2 |
| 75 | + |
| 76 | + # Forward Euler step |
| 77 | + u_next = u + dt * dudt |
| 78 | + |
| 79 | + # Enforce boundary conditions (Dirichlet: hold boundary values) |
| 80 | + u_next = torch.cat([u[:1], u_next[1:-1], u[-1:]]) |
| 81 | + |
| 82 | + return {"u_next": u_next} |
| 83 | + |
| 84 | + |
| 85 | +def apply(inputs: InputSchema) -> OutputSchema: |
| 86 | + tensor_inputs = tree_map(to_tensor, inputs.model_dump()) |
| 87 | + return evaluate(tensor_inputs) |
| 88 | + |
| 89 | + |
| 90 | +def abstract_eval(abstract_inputs: Any) -> Any: |
| 91 | + return {"u_next": {"shape": [N], "dtype": "float64"}} |
| 92 | + |
| 93 | + |
| 94 | +def jacobian_vector_product( |
| 95 | + inputs: InputSchema, |
| 96 | + jvp_inputs: set[str], |
| 97 | + jvp_outputs: set[str], |
| 98 | + tangent_vector: dict[str, Any], |
| 99 | +): |
| 100 | + jvp_inputs = tuple(jvp_inputs) |
| 101 | + tangent_vector = {key: tangent_vector[key] for key in jvp_inputs} |
| 102 | + |
| 103 | + tensor_inputs = tree_map(to_tensor, inputs.model_dump()) |
| 104 | + pos_tangent = tree_map(to_tensor, tangent_vector).values() |
| 105 | + pos_inputs = flatten_with_paths(tensor_inputs, jvp_inputs).values() |
| 106 | + |
| 107 | + filtered_pos_eval = filter_func( |
| 108 | + evaluate, tensor_inputs, jvp_outputs, input_paths=jvp_inputs |
| 109 | + ) |
| 110 | + |
| 111 | + return torch.func.jvp(filtered_pos_eval, tuple(pos_inputs), tuple(pos_tangent))[1] |
| 112 | + |
| 113 | + |
| 114 | +def vector_jacobian_product( |
| 115 | + inputs: InputSchema, |
| 116 | + vjp_inputs: set[str], |
| 117 | + vjp_outputs: set[str], |
| 118 | + cotangent_vector: dict[str, Any], |
| 119 | +): |
| 120 | + vjp_inputs = tuple(vjp_inputs) |
| 121 | + cotangent_vector = {key: cotangent_vector[key] for key in vjp_outputs} |
| 122 | + |
| 123 | + tensor_inputs = tree_map(to_tensor, inputs.model_dump()) |
| 124 | + tensor_cotangent = tree_map(to_tensor, cotangent_vector) |
| 125 | + pos_inputs = flatten_with_paths(tensor_inputs, vjp_inputs).values() |
| 126 | + |
| 127 | + filtered_pos_func = filter_func( |
| 128 | + evaluate, tensor_inputs, vjp_outputs, input_paths=vjp_inputs |
| 129 | + ) |
| 130 | + |
| 131 | + _, vjp_func = torch.func.vjp(filtered_pos_func, *pos_inputs) |
| 132 | + vjp_vals = vjp_func(tensor_cotangent) |
| 133 | + return dict(zip(vjp_inputs, vjp_vals, strict=True)) |
| 134 | + |
| 135 | + |
| 136 | +def jacobian( |
| 137 | + inputs: InputSchema, |
| 138 | + jac_inputs: set[str], |
| 139 | + jac_outputs: set[str], |
| 140 | +): |
| 141 | + jac_inputs = tuple(jac_inputs) |
| 142 | + tensor_inputs = tree_map(to_tensor, inputs.model_dump()) |
| 143 | + pos_inputs = flatten_with_paths(tensor_inputs, jac_inputs).values() |
| 144 | + |
| 145 | + filtered_pos_eval = filter_func( |
| 146 | + evaluate, tensor_inputs, jac_outputs, input_paths=jac_inputs |
| 147 | + ) |
| 148 | + |
| 149 | + def filtered_pos_eval_flat(*args): |
| 150 | + res = filtered_pos_eval(*args) |
| 151 | + return tuple(res[k] for k in jac_outputs) |
| 152 | + |
| 153 | + jac = torch.autograd.functional.jacobian(filtered_pos_eval_flat, tuple(pos_inputs)) |
| 154 | + |
| 155 | + jac_dict = {} |
| 156 | + for dy, dys in zip(jac_outputs, jac, strict=True): |
| 157 | + jac_dict[dy] = {} |
| 158 | + for dx, dxs in zip(jac_inputs, dys, strict=True): |
| 159 | + jac_dict[dy][dx] = dxs |
| 160 | + |
| 161 | + return jac_dict |
0 commit comments