def evaluate_primal(
firedrake_function: Callable[..., BackendVariable],
firedrake_templates: Collection[BackendVariable],
*args: np.array,
) -> Tuple[np.array, BackendVariable, Collection[BackendVariable], pyadjoint.Tape]:
"""Computes the output of a firedrake_function and saves a corresponding pyadjoint tape
Input:
firedrake_function (callable): Firedrake function to be executed during the forward pass
firedrake_templates (collection of BackendVariable): Templates for converting arrays to Firedrake types
args (tuple): NumPy array representation of the input to firedrake_function
Output:
numpy_output (np.array): NumPy array representation of the output from firedrake_function(*firedrake_inputs)
firedrake_output (AdjFloat or Function): Firedrake representation of the output from firedrake_function(*firedrake_inputs)
firedrake_inputs (collection of BackendVariable): Firedrake representation of the input args
tape (pyadjoint.Tape): pyadjoint's saved computational graph
"""
check_input(firedrake_templates, *args)
firedrake_inputs = convert_all_to_backend(firedrake_templates, *args)
# Create tape associated with this forward pass
tape = pyadjoint.Tape()
pyadjoint.set_working_tape(tape)
firedrake_output = firedrake_function(*firedrake_inputs)
if isinstance(firedrake_output, tuple):
raise ValueError("Only single output from Firedrake function is supported.")
numpy_output = np.asarray(to_numpy(firedrake_output))
return numpy_output, firedrake_output, firedrake_inputs, tape
def evaluate_primal(
fenics_function: Callable,
fenics_templates: Iterable[FenicsVariable],
*args: np.array,
) -> Tuple[np.array, FenicsVariable, Tuple[FenicsVariable], pyadjoint.Tape]:
"""Computes the output of a fenics_function and saves a corresponding gradient tape
Input:
fenics_function (callable): FEniCS function to be executed during the forward pass
fenics_templates (iterable of FenicsVariable): Templates for converting arrays to FEniCS types
args (tuple): NumPy array representation of the input to fenics_function
Output:
numpy_output (np.array): NumPy array representation of the output from fenics_function(*fenics_inputs)
fenics_output (AdjFloat or Function): FEniCS representation of the output from fenics_function(*fenics_inputs)
fenics_inputs (list of FenicsVariable): FEniCS representation of the input args
tape (pyadjoint.Tape): pyadjoint's saved computational graph
"""
check_input(fenics_templates, *args)
fenics_inputs = convert_all_to_fenics(fenics_templates, *args)
# Create tape associated with this forward pass
tape = pyadjoint.Tape()
pyadjoint.set_working_tape(tape)
fenics_output = fenics_function(*fenics_inputs)
if isinstance(fenics_output, tuple):
raise ValueError(
"Only single output from FEniCS function is supported.")
numpy_output = np.asarray(fenics_to_numpy(fenics_output))
return numpy_output, fenics_output, fenics_inputs, tape
def fem_eval(
fenics_function: Callable,
fenics_templates: Iterable[FenicsVariable],
*args: np.array,
) -> Tuple[np.array, FenicsVariable, Tuple[FenicsVariable], pyadjoint.Tape]:
"""Computes the output of a fenics_function and saves a corresponding gradient tape
Input:
fenics_function (callable): FEniCS function to be executed during the forward pass
fenics_templates (iterable of FenicsVariable): Templates for converting arrays to FEniCS types
args (tuple): jax array representation of the input to fenics_function
Output:
numpy_output (np.array): JAX array representation of the output from fenics_function(*fenics_inputs)
residual_form (ufl.Form): UFL Form for the residual used to solve the problem with fenics.solve(F==0, ...)
fenics_inputs (list of FenicsVariable): FEniCS representation of the input args
"""
check_input(fenics_templates, *args)
fenics_inputs = convert_all_to_fenics(fenics_templates, *args)
# Create tape associated with this forward pass
tape = pyadjoint.Tape()
pyadjoint.set_working_tape(tape)
fenics_output = fenics_function(*fenics_inputs)
if isinstance(fenics_output, tuple):
raise ValueError(
"Only single output from FEniCS function is supported.")
numpy_output = np.asarray(fenics_to_numpy(fenics_output))
return numpy_output, fenics_output, fenics_inputs, tape
def pytest_runtest_setup(item):
""" Hook function which is called before every test """
set_working_tape(Tape())
# Fix the seed to avoid random test failures due to slight tolerance variations
numpy.random.seed(21)
from .assembly import assemble, assemble_system
from .solving import solve
from .projection import project
from .interpolation import interpolate
from .ufl_constraints import UFLEqualityConstraint, UFLInequalityConstraint
from .shapead_transformations import (transfer_from_boundary,
transfer_to_boundary)
if backend.__name__ != "firedrake":
from .newton_solver import NewtonSolver
from .lu_solver import LUSolver
from .krylov_solver import KrylovSolver
from .petsc_krylov_solver import PETScKrylovSolver
from .types import *
from .refine import refine
from .system_assembly import *
from .variational_solver import (NonlinearVariationalProblem,
NonlinearVariationalSolver,
LinearVariationalProblem,
LinearVariationalSolver)
from pyadjoint import (Tape, set_working_tape, get_working_tape,
pause_annotation, continue_annotation,
ReducedFunctional, taylor_test, taylor_to_dict,
compute_gradient, compute_hessian, AdjFloat, Control,
minimize, maximize, MinimizationProblem, IPOPTSolver,
ROLSolver, InequalityConstraint, EqualityConstraint,
MoolaOptimizationProblem, print_optimization_methods,
stop_annotating)
set_working_tape(Tape())