Gradient Optimization Circuit Search Jacobians Stream Calculus Rewriting Multi-Dim Streams Coupled SDEs JAX & JIT Engineering Examples

Multi-Dimensional Streams

Work with 2D and 3D data for PDEs and spatial problems.

Shape Conventions

A Stream has shape (chunk_size, dim1, dim2, ...):

import jax.numpy as jnp
from gimle.asgard.runtime.stream import Stream

# Scalar stream (no spatial dimensions)
Stream(data=jnp.array([1.0, 2.0, 3.0]), dim_labels=(), chunk_size=3)
# shape: (3,)

# 1D spatial (5 coefficients along x)
Stream(data=jnp.array([[1.0, 2.0, 3.0, 4.0, 5.0]]),
       dim_labels=("x",), chunk_size=1)
# shape: (1, 5)

# 2D spatial (4x4 grid)
Stream(data=jnp.zeros((1, 4, 4)),
       dim_labels=("x", "y"), chunk_size=1)
# shape: (1, 4, 4)

Dimension labels must be unique. The order matters — it determines which axis corresponds to which label.

Creating 2D Input Data

Python API

from gimle.asgard.api import GimleAPI
import jax.numpy as jnp

gimle = GimleAPI()

# 2D coefficient array: shape (1, rows, cols)
data = jnp.array([[[1.0, 2.0, 1.0],
                    [3.0, 4.0, 2.0],
                    [1.0, 1.0, 0.0]]])

eq = gimle.create_equation("diff(f, x) = g")
circuit = gimle.compile(eq)
result = gimle.simulate(circuit, [data], dim_labels=("x", "y"))

YAML Examples

input:
  shape: [3, 3]
  values:
    - [1.0, 2.0, 1.0]
    - [3.0, 4.0, 2.0]
    - [1.0, 1.0, 0.0]
  dims: [x, y]

evaluate:
  x:
    range: [-1, 1]
    points: 30
  y:
    range: [-1, 1]
    points: 30

output:
  type: Solution2D
  title: "2D Solution"
  labels:
    x: x
    y: y
    z: f(x,y)

Evaluation Across Dimensions

When you evaluate a multi-dimensional stream at specific points, the evaluator contracts each dimension using basis functions from its calculus:

from gimle.asgard.runtime.stream_evaluator import StreamEvaluator, RealCalculus

evaluator = StreamEvaluator(stream, {
    "x": RealCalculus(center=0.0),
    "y": RealCalculus(center=0.0),
})

# Single point
value = evaluator.evaluate(x=1.0, y=2.0)

# Grid evaluation (30x30 points)
values = evaluator.evaluate(
    x=jnp.linspace(-1, 1, 30),
    y=jnp.linspace(-1, 1, 30),
)

Each dimension is contracted independently using Taylor basis evaluation: basis(degree, point) = (point - center)^degree / degree!.

Register and Deregister on Multi-Dimensional Streams

Register and deregister operate on a single named dimension:

# Register along x: shifts rows down, prepends zero row
# [[1, 2, 3],     [[0, 0, 0],
#  [4, 5, 6]]  →   [1, 2, 3]]

# Register along y: shifts columns right, prepends zero column
# [[1, 2, 3],     [[0, 1, 2],
#  [4, 5, 6]]  →   [0, 4, 5]]

The dimension name comes from the circuit syntax: register(x) operates on the x dimension, register(y) on the y dimension.

2D PDE Examples

Heat Equation Analogue

name: Heat Equation Analogue
equation: "diff(f, x) = diff(diff(f, y), y)"

input:
  shape: [5, 6]
  fill: 1.0
  dims: [x, y]

evaluate:
  x:
    range: [-1, 1]
    points: 30
  y:
    range: [-1, 1]
    points: 30

output:
  type: Solution2D
  title: "df/dx = d²f/dy²"

Mixed Partial Derivatives

name: Mixed Derivatives
equation: "diff(diff(f, x), y) = g"

input:
  shape: [3, 3]
  values:
    - [1.0, 2.0, 1.0]
    - [3.0, 4.0, 2.0]
    - [1.0, 1.0, 0.0]
  dims: [x, y]

evaluate:
  x:
    range: [-1, 1]
    points: 30
  y:
    range: [-1, 1]
    points: 30

output:
  type: Solution2D
  title: "d²f/dxdy"

Limitations