import jax.numpy as np
import jax
import os
import glob
import jaxopt
from jax_am.fem.generate_mesh import Mesh
from jax_am.fem.solver import ad_wrapper
from jax_am.fem.utils import save_sol
from jax_am.common import rectangle_mesh
from fem_model import Elasticity
problem_name = 'plate'
root_path = os.path.join(os.path.dirname(__file__), 'data')
files = glob.glob(os.path.join(root_path, f'vtk/{problem_name}/*'))
for f in files:
os.remove(f)
L = 60.
W = 30.
N_L = 60
N_W = 30
meshio_mesh = rectangle_mesh(N_L, N_W, L, W)
jax_mesh = Mesh(meshio_mesh.points, meshio_mesh.cells_dict['quad'])
def fixed_location(point):
return np.isclose(point[0], 0., atol=1e-5)
def load_location(point):
return np.logical_and(np.isclose(point[0], L, atol=1e-5),
np.isclose(point[1], 0., atol=1. + 1e-5))
def dirichlet_val(point):
return 0.
def neumann_val(point):
return np.array([0., -1])
dirichlet_bc_info = [[fixed_location] * 2, [0, 1], [dirichlet_val] * 2]
neumann_bc_info = [[load_location], [neumann_val]]
problem = Elasticity(jax_mesh,
vec=2,
dim=2,
ele_type='QUAD4',
dirichlet_bc_info=dirichlet_bc_info,
neumann_bc_info=neumann_bc_info,
additional_info=(problem_name, ))
fwd_pred = ad_wrapper(problem, linear=True, use_petsc=True)
def J_fn(dofs, params):
"""J(u, p)
"""
sol = dofs.reshape((problem.num_total_nodes, problem.vec))
compliance = problem.compute_compliance(neumann_val, sol)
return compliance
def J_total(params):
"""J(u(p), p)
"""
sol = fwd_pred(params)
dofs = sol.reshape(-1)
obj_val = J_fn(dofs, params)
return obj_val
outputs = []
def output_sol(params, obj_val):
print("\nOutput solution - need to solve the forward problem again...")
sol = fwd_pred(params)
vtu_path = os.path.join(
root_path, f'vtk/{problem_name}/sol_{output_sol.counter:03d}.vtu')
save_sol(problem,
sol,
vtu_path,
cell_infos=[('theta', problem.full_params[:, 0])])
print(f"compliance = {obj_val}")
outputs.append(obj_val)
output_sol.counter += 1
output_sol.counter = 0
vf = 0.5
def objectiveHandle(rho):
J, dJ = jax.value_and_grad(J_total)(rho)
output_sol(rho, J)
return J, dJ
# Initialize params
params = vf * np.ones((len(problem.flex_inds), 1))
print(f'Initial compliance is: {J_total(params):.5e}')
solver = jaxopt.LBFGS(fun=objectiveHandle, value_and_grad=True, has_aux=False)
# Trainiting loop
opt_state = solver.init_state(params)
for _ in range(2):
params, opt_state = solver.update(params=params, state=opt_state)
Traceback (most recent call last):
File "/home/igork/IRP/NeurOpt/jax-am-forked/applications/fem/top_opt/jaxopt_mre.py", line 114, in <module>
params, opt_state = solver.update(params=params, state=opt_state)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jaxopt/_src/lbfgs.py", line 348, in update
ls_state = zoom_linesearch(f=self._value_and_grad_with_aux,
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jaxopt/_src/zoom_linesearch.py", line 467, in zoom_linesearch
state = lax.while_loop(lambda state: (~state.done) & (state.i <= maxiter) & (~state.failed),
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/lax/control_flow/loops.py", line 1150, in while_loop
init_vals, init_avals, body_jaxpr, in_tree, *rest = _create_jaxpr(init_val)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/lax/control_flow/loops.py", line 1133, in _create_jaxpr
body_jaxpr, body_consts, body_tree = _initial_style_jaxpr(
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/lax/control_flow/common.py", line 60, in _initial_style_jaxpr
jaxpr, consts, out_tree = _initial_style_open_jaxpr(
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/lax/control_flow/common.py", line 54, in _initial_style_open_jaxpr
jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals, debug)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2150, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2172, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers_)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/linear_util.py", line 188, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jaxopt/_src/zoom_linesearch.py", line 385, in body
(phi_i, dphi_i, g_i), aux_i = restricted_func_and_grad(a_i)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jaxopt/_src/zoom_linesearch.py", line 329, in restricted_func_and_grad
(phi, aux), g = f_value_and_grad(xkp1, *args, **kwargs)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jaxopt/_src/base.py", line 62, in value_and_grad_with_aux
v, g = value_and_grad(*a, **kw)
File "/home/igork/IRP/NeurOpt/jax-am-forked/applications/fem/top_opt/jaxopt_mre.py", line 99, in objectiveHandle
J, dJ = jax.value_and_grad(J_total)(rho)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/api.py", line 718, in value_and_grad_f
ans, vjp_py = _vjp(f_partial, *dyn_args, reduce_axes=reduce_axes)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/api.py", line 2174, in _vjp
out_primal, out_vjp = ad.vjp(
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/interpreters/ad.py", line 139, in vjp
out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/interpreters/ad.py", line 128, in linearize
jaxpr, out_pvals, consts = pe.trace_to_jaxpr_nounits(jvpfun_flat, in_pvals)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 777, in trace_to_jaxpr_nounits
jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/linear_util.py", line 188, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/igork/IRP/NeurOpt/jax-am-forked/applications/fem/top_opt/jaxopt_mre.py", line 70, in J_total
sol = fwd_pred(params)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/custom_derivatives.py", line 614, in __call__
out_flat = custom_vjp_call_p.bind(flat_fun, flat_fwd, flat_bwd,
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/custom_derivatives.py", line 763, in bind
outs = top_trace.process_custom_vjp_call(self, fun, fwd, bwd_, tracers,
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/interpreters/ad.py", line 396, in process_custom_vjp_call
res_and_primals_out = fwd.call_wrapped(*fwd_in)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/linear_util.py", line 188, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax_am/fem/solver.py", line 545, in f_fwd
sol = fwd_pred(params)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/custom_derivatives.py", line 614, in __call__
out_flat = custom_vjp_call_p.bind(flat_fun, flat_fwd, flat_bwd,
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/custom_derivatives.py", line 763, in bind
outs = top_trace.process_custom_vjp_call(self, fun, fwd, bwd_, tracers,
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 1985, in process_custom_vjp_call
fun_jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, self.main, in_avals)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2172, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers_)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/linear_util.py", line 188, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax_am/fem/solver.py", line 541, in fwd_pred
sol = solver(problem, linear=linear, use_petsc=use_petsc)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax_am/fem/solver.py", line 473, in solver
return solver_row_elimination(problem, linear, precond, initial_guess, use_petsc)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax_am/fem/solver.py", line 228, in solver_row_elimination
res_vec, A_fn = newton_update_helper(dofs)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax_am/fem/solver.py", line 220, in newton_update_helper
res_vec = problem.newton_update(dofs.reshape(sol_shape)).reshape(-1)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax_am/fem/core.py", line 665, in newton_update
return self.compute_newton_vars(sol, **self.internal_vars)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax_am/fem/core.py", line 640, in compute_newton_vars
weak_form, cells_jac = self.split_and_compute_cell(cells_sol, onp, True, **internal_vars)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax_am/fem/core.py", line 519, in split_and_compute_cell
values = np_version.vstack(values)
File "<__array_function__ internals>", line 200, in vstack
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/numpy/core/shape_base.py", line 293, in vstack
arrs = atleast_2d(*tup)
File "<__array_function__ internals>", line 200, in atleast_2d
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/numpy/core/shape_base.py", line 121, in atleast_2d
ary = asanyarray(ary)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax/_src/core.py", line 598, in __array__
raise TracerArrayConversionError(self)
jax._src.traceback_util.UnfilteredStackTrace: jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(float64[90,4,2])>with<DynamicJaxprTrace(level=1/1)>
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError
The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.
--------------------
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/home/igork/IRP/NeurOpt/jax-am-forked/applications/fem/top_opt/jaxopt_mre.py", line 114, in <module>
params, opt_state = solver.update(params=params, state=opt_state)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jaxopt/_src/lbfgs.py", line 348, in update
ls_state = zoom_linesearch(f=self._value_and_grad_with_aux,
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jaxopt/_src/zoom_linesearch.py", line 467, in zoom_linesearch
state = lax.while_loop(lambda state: (~state.done) & (state.i <= maxiter) & (~state.failed),
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jaxopt/_src/zoom_linesearch.py", line 385, in body
(phi_i, dphi_i, g_i), aux_i = restricted_func_and_grad(a_i)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jaxopt/_src/zoom_linesearch.py", line 329, in restricted_func_and_grad
(phi, aux), g = f_value_and_grad(xkp1, *args, **kwargs)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jaxopt/_src/base.py", line 62, in value_and_grad_with_aux
v, g = value_and_grad(*a, **kw)
File "/home/igork/IRP/NeurOpt/jax-am-forked/applications/fem/top_opt/jaxopt_mre.py", line 99, in objectiveHandle
J, dJ = jax.value_and_grad(J_total)(rho)
File "/home/igork/IRP/NeurOpt/jax-am-forked/applications/fem/top_opt/jaxopt_mre.py", line 70, in J_total
sol = fwd_pred(params)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax_am/fem/solver.py", line 545, in f_fwd
sol = fwd_pred(params)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax_am/fem/solver.py", line 541, in fwd_pred
sol = solver(problem, linear=linear, use_petsc=use_petsc)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax_am/fem/solver.py", line 473, in solver
return solver_row_elimination(problem, linear, precond, initial_guess, use_petsc)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax_am/fem/solver.py", line 228, in solver_row_elimination
res_vec, A_fn = newton_update_helper(dofs)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax_am/fem/solver.py", line 220, in newton_update_helper
res_vec = problem.newton_update(dofs.reshape(sol_shape)).reshape(-1)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax_am/fem/core.py", line 665, in newton_update
return self.compute_newton_vars(sol, **self.internal_vars)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax_am/fem/core.py", line 640, in compute_newton_vars
weak_form, cells_jac = self.split_and_compute_cell(cells_sol, onp, True, **internal_vars)
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/jax_am/fem/core.py", line 519, in split_and_compute_cell
values = np_version.vstack(values)
File "<__array_function__ internals>", line 200, in vstack
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/numpy/core/shape_base.py", line 293, in vstack
arrs = atleast_2d(*tup)
File "<__array_function__ internals>", line 200, in atleast_2d
File "/home/igork/miniconda3/envs/jax-nrto/lib/python3.10/site-packages/numpy/core/shape_base.py", line 121, in atleast_2d
ary = asanyarray(ary)
jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(float64[90,4,2])>with<DynamicJaxprTrace(level=1/1)>
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError
So it seems that the incompatibility arises in jax/fem/core.py. Have you encountered a similar error before or do you have any ideas how it could be solved? I'm suspecting that the issue might have to do with the fact that jaxopt introduces its own decorators for implicit differentiation and they might clash with ad_wrapper.