pdesolve

pdesolve integrates a problem with a solver over a time span and returns its Solution, and compile compiles that call without running it. See running a simulation.

Recommended defaults

Most calls pass a problem, a solver, the span, and the initial state, plus saveat= when the final state alone is not enough. Leave adjoint=None unless you need jax.grad through the solve. Leave lambda0, consistent_ic, and control_state0 unset unless a DAE's initialization or an operating protocol demands them.

diff_soup.pdesolve.pdesolve

pdesolve(problem: Any, solver: Any, t_span: tuple, y0: Any, *, saveat: Any = None, lambda0: Optional[Any] = None, theta: Any = None, adjoint: Optional[Any] = None, event: Optional[Any] = None, control_state0: Optional[Any] = None, consistent_ic: Optional[bool] = None, trace: bool = False, **solver_kwargs: Any) -> 'Solution'

Integrate a problem over a time span and return its Solution.

Every keyword reaches solver.solve, so engine keywords such as dt0, max_steps and tolerances pass through. A one-time warning fires when JAX's 64-bit mode is off.

Parameters:

Name	Type	Description	Default
problem	Any	A LoweredProblem, MassProblem, or ImplicitProblem.	required
solver	Any	A solver with a solve method, such as JNDF, JBDF or JIDA.	required
t_span	tuple	Start and end times (t0, t1).	required
y0	Any	Initial state, the full state vector or, when lambda0 is given, the field block.	required
saveat	Any	A SaveAt choosing the saved outputs, or None for the final state only.	None
lambda0	Optional[Any]	Initial values of the global unknowns, packed after y0.	None
theta	Any	Parameter values, or None for the problem's declared values.	None
adjoint	Optional[Any]	Differentiation strategy, such as CheckpointAdjoint for reverse mode, or None for the solver's own.	None
event	Optional[Any]	An Event or EventSet monitored during the solve.	None
control_state0	Optional[Any]	Initial operating point of a problem with declared controls.	None
consistent_ic	Optional[bool]	Whether to make the initial state consistent before the first step, or None for the solver's default.	None
trace	bool	If True, attach a per-step diagnostic payload to the solution.	False
**solver_kwargs	Any	Engine keywords for solver.solve, such as dt0, max_steps, tolerances and linear.	{}

Returns:

Type	Description
'Solution'	The Solution of the solve.

diff_soup.compile

compile(problem: Any, solver: Any, t_span: tuple, y0: Any, *, saveat: Any = None, theta: Any = None, lambda0: Optional[Any] = None, adjoint: Optional[Any] = None, event: Optional[Any] = None, control_state0: Optional[Any] = None, consistent_ic: Optional[bool] = None, **solver_kwargs: Any) -> CompiledSolve

Compile a pdesolve call without running it.

The arguments are those of pdesolve. t_span, y0, theta and lambda0 fix only shapes and dtypes, and every other argument, the save times included, is fixed in the compiled solve. Nothing is solved here.

compiled = ds.compile(problem, solver, (0.0, 0.1), y0)
sol = compiled((0.0, 0.2), 2 * y0)

Parameters:

Name	Type	Description	Default
problem	Any	The problem, as for pdesolve.	required
solver	Any	The solver, as for pdesolve.	required
t_span	tuple	Start and end times (t0, t1), whose values may change at each call.	required
y0	Any	Initial state, whose values may change at each call.	required
saveat	Any	A SaveAt, whose ts must be ordered=True.	None
theta	Any	Parameter values, then required at each call, or None for the problem's declared values.	None
lambda0	Optional[Any]	Initial values of the global unknowns, then required at each call.	None
adjoint	Optional[Any]	Differentiation strategy, as for pdesolve.	None
event	Optional[Any]	An Event monitored during the solve.	None
control_state0	Optional[Any]	Initial operating point of a problem with declared controls.	None
consistent_ic	Optional[bool]	Whether to make the initial state consistent before the first step.	None
**solver_kwargs	Any	Engine keywords for solver.solve, as for pdesolve.	{}

Returns:

Type	Description
CompiledSolve	A CompiledSolve that runs the solve.

Raises:

Type	Description
ValueError	When saveat.ts is given with ordered=False.

CompiledSolve

A solve compiled by compile, ready to run.

Call it with new values of the shapes and dtypes it was compiled for. Every other argument of the solve is fixed at compile time.

Attributes:

Name	Type	Description
compile_seconds		The wall time of tracing and compiling the solve, in seconds.

__call__

__call__(t_span: tuple, y0: Any, *, theta: Any = None, lambda0: Any = None)

Run the compiled solve from new values.

Parameters:

Name	Type	Description	Default
t_span	tuple	Start and end times (t0, t1).	required
y0	Any	Initial state.	required
theta	Any	Parameter values, required when compile was given them and None otherwise.	None
lambda0	Any	Initial values of the global unknowns, required when compile was given them and None otherwise.	None

Returns:

Type	Description
	The Solution of the solve.

Raises:

Type	Description
ValueError	When an argument's shape or dtype differs from the compiled one.

cost_analysis

cost_analysis() -> Any

The compiled executable's cost estimate, as jax.stages.Compiled.cost_analysis reports it.

memory_analysis

memory_analysis() -> Any

The compiled executable's memory use, as jax.stages.Compiled.memory_analysis reports it.