tmdsimpy.jax.NonlinearSolverOMP

class tmdsimpy.jax.NonlinearSolverOMP(config={})

Bases: NonlinearSolver

Parallel nonlinear solver object that contains several functions.

This solver instance utilizes JAX for shared memory parallelism of linear algebra steps for the linear and nonlinear solves.

Parameters:

configdict, optional

Dictionary of settings to be used in the solver. All keys are optional. Keys include:

max_stepsint, default 20

Maximum number of iterations allowed in the nonlinear solver.

reform_freqint, default 1

Frequency of recalculating and refactoring Jacobian matrix of the nonlinear problem. 1 corresponds to Newton-Raphson (every step). Larger numbers correspond to BFGS low rank updates in between steps with refactoring. When reform_freq > 1, function being solved must accept the keyword calc_grad=True or calc_grad=False to differentiate if Jacobian matrix should be calculated. If calc_grad=True, then returned tuple should be (R, dRdX). If False, returned tuple should start with (R,), but may return other values past the 0th index of tuple. Function may be a lambda function that completely ignores calc_grad, for instance: fun = lambda X, calc_grad=True : fun0(X)

verbosebool, default True

Flag for if output should be printed.

xtoldouble, default None

Convergence tolerance on the L2 norm of the step size (dX) for nsolve. If None, code will set the value to be equal to 1e-6*X0.shape[0] where X0 is the initial guess for a given solution calculation. if xtol is passed to nsolve, that value is used instead

rtoldouble, default None

Convergence toleranace on the L2 norm of the residual vector (R) for nsolve.

etoldouble, default None

Convergence tolerance on the energy norm of the inner product of step (dX) and residual (R) or e=np.abs(dX @ R) for nsolve.

xtol_reldouble, default None

Convergence tolerance on norm(dX) / norm(dX_step0) for nsolve.

rtol_reldouble, default None

Convergence tolerance on norm(R) / norm(R_step0) for nsolve.

etol_reldouble, default None

Convergence tolerance on norm(e) / norm(e_step0) for nsolve.

stopping_tol: list, default [‘xtol’]

List can contain options of ‘xtol’, ‘rtol’, ‘etol’, ‘xtol_rel’, ‘rtol_rel’, ‘etol_rel’. If any of the listed tolerances are satisfied, then iteration is considered converged and exits. Futher development would allow for the list to contain lists of these same options and in a sublist, all options would be required. This has not been implemented.

accepting_tollist, default []

List that can contain the same set of strings as stopping_tol. Once maximum interactions has been reached, if any of these tolerances are satisified by the final step, then the solution is considered converged. This allows for looser tolerances to be accepted instead of non-convergence, while still using max iterations to try to achieve the tighter tolerances.

line_search_itersint, default 0

Number of iterations used in line_search. If 0, line search is not used. If line search is desired, a recommended value is less than 10, perhaps about 2 to 5. If it is greater than 0, then function being solved must accept calc_grad=True or calc_grad=False as inputs. The default is 0.

line_search_tolfloat, default 0.5

If the line search function decreases to be less than the initial value times this tolerance, than it is accepted as converged. This is not intended to be a tight tolerance, line search is just intended to quickly reduce the step in case of poor problem behavior. The default is 0.5.

line_search_same_signbool, default True

If true, line search tries to only a return a step that does not change the sign of the quantity deltaX @ R(X + alpha*deltaX). The default is True.

See also

tmdsimpy.NonlinearSolver

Base class that supports similar operations with less parallelism.

nsolve

Nonlinear solution method on this class.

Notes

This class inherets some members from tmdsimpy.NonlinearSolver. For instance, eigs could be substituted to try to have better parallelism here.

__init__(config={})

Methods

__init__([config])
conditioning_wrapper(fun, CtoP[, RPtoC])	Function to create a wrap a provided function and add conditioning.
edit_config(new_config)	Edit the solver configuration settings after initialization.
eigs(K[, M, subset_by_index])	Conduct eigenvalue analysis for a linear system.
lin_factor(A)	Do an LU factorization of a matrix for later linear solutions.
lin_factored_solve(lu_and_piv, b)	Solve a linear system that has already been factored.
lin_solve(A, b)	Solve the linear system A @ x = b
line_search(fun, X, Rx, deltaX)	Line search algorithm to help in the numerical solution to a set of nonlinear equations.
nsolve(fun, X0[, verbose, xtol, Dscale])	Numerically solves multiple nonlinear equations to find roots.

conditioning_wrapper(fun, CtoP, RPtoC=1.0)

Function to create a wrap a provided function and add conditioning.

Parameters:

funfunction

Function that is to be wrapped in condition space. Function should take two arguments, one is unknown vector Xp, other is optional argument of calc_grad=True. The calc_grad=True is only passed to fun, if calc_grad=False is passed to the fun_conditioned that is returned.

CtoP(N,) numpy.ndarray

Vector describing conversion from physical coordinates to conditioned coordinates for the unknown vector that is input to fun.

RPtoCfloat, optional

Scales the full ouptut residual vector by this magnitude. The default is 1.0.

Returns:

fun_conditionedfunction

Function that describes the same nonlinear problem as fun, but in a conditioned space. Function takes input of Xc where Xp = CtoP * Xc. Second optional input to the fuction is calc_grad=True. This function returns a residual vector for conditioned inputs. If calc_grad=True, the Jacobian in conditioned space is also returned by this function.

edit_config(new_config)

Edit the solver configuration settings after initialization.

Parameters:

new_configdict

Dictionary of key value pairs for the new settings. The settings that can be changed are the same as those when creating a new object.

Returns:

None.

See also

NonlinearSolverOMP

Class documentation for the setting keys that can be changed.

eigs(K, M=None, subset_by_index=[0, 2])

Conduct eigenvalue analysis for a linear system.

Parameters:

K(N,N) numpy.ndarray

Stiffness matrix for general second order system. Assumed to be symmetric.

M(N,N) numpy.ndarray, optional

Mass matrix for general second order system. This must be symmetric positive definite. The default is None.

subset_by_indexlist of length 2, optional

Subset indices for which eigenvalues should be calculated. See scipy.linalg.eigh for more details. Let M be the number of requested eigenvalues by this parameter by setting it as [0,M]. The default is [0, 2].

Returns:

eigvals(M,) numpy.ndarray

Eigenvalues of the linear problem.

eigvecs(N,M) numpy.ndarray

Eigenvectors of linear problem with columns corresponding to individual eigenvalues.

See also

scipy.linalg.eigh

The eigen problem solver that is called here. This just provides an interface that can be made consistent with different approaches.

lin_factor(A)

Do an LU factorization of a matrix for later linear solutions.

Parameters:

A(N,N) numpy.ndarray

Linear system matrix for later solving.

Returns:

lu_and_pivtuple

Resulting data from factoring the matrix A, can be passed to lin_factored_solve to solve the linear system.

See also

lin_factored_solve

Method for solving after factoring.

Notes

This method calls jax.scipy.linalg.lu_factor(A).

lin_factored_solve(lu_and_piv, b)

Solve a linear system that has already been factored.

Parameters:

lu_and_pivtuple

Results from factoring a matrix with lin_factor(A)

b(N,) numpy.ndarray

Right hand side vector.

Returns:

x(N,) numpy.ndarray

Solution to the linear problem

See also

lin_factor

Method for factoring matrix prior to this solution.

lin_solve(A, b)

Solve the linear system A @ x = b

Parameters:

A(N,N) numpy.ndarray

Linear system matrix.

b(N,) numpy.ndarray

Right hand side vector.

Returns:

x(N,) numpy.ndarray

Solution to the linear problem

Notes

Implementation uses jax.numpy.linalg.solve, which supports some parallelism.

line_search(fun, X, Rx, deltaX)

Line search algorithm to help in the numerical solution to a set of nonlinear equations. This is used by nsolve.

Parameters:

funfunction

Function to be solved, returns R=fun(X, calc_grad={True or False})[0]. Must accept the input argument calc_grad=True and calc_grad=False. Function may be a lambda function that completely ignores calc_grad.

X(N,) numpy.ndarray

Values of unknowns at initial point.

Rx(N,) numpy.ndarray

Residual at X, equal to fun(X)[0].

deltaX(N,) numpy.ndarray

Step direction of interest, generally calculated based on gradient solution step.

Returns:

alphafloat

Fraction of deltaX step that should be taken. Recommended update is X = X + alpha*deltaX

soldict

Description of final solution state. Has keys of [‘message’, ‘nfev’]. ‘nfev’ is the number of function evaluations completed. ‘message’ describes how line search exited.

See also

NonlinearSolverOMP

Documentation for the solver class describes configurations and settings of the numerical solver that are configured at creation rather than at solution time. Relevant settings are line_search_iters, line_search_tol, line_search_same_sign.

nsolve

Nonlinear solver routine that calls this function (Newton-Raphson + BFGS Solver)

Notes

The line search algorithm here is based on [1]. In the finite element context, the objective is to find the zero of an ‘energy’ norm of R^T deltaX. The bisection algorithm is used to find a solution for R(X + alpha*deltaX) with alpha in [0, 1]. A very loose tolerance is generally desired to minimize additional computational cost of this function before returing to nsolve.

References

[1]

Matthies, H., G. Strang, 1979. The solution of nonlinear finite element equations. International Journal for Numerical Methods in Engineering 14, 1613–1626. https://doi.org/10.1002/nme.1620141104

nsolve(fun, X0, verbose=None, xtol=None, Dscale=1.0)

Numerically solves multiple nonlinear equations to find roots.

Solver settings are set at initialization of NonlinearSolverOMP (see that documentation).

Parameters:

funfunction handle

Function to be solved, function returns two arguments of R (residual, (N,) numpy.ndarray) and dRdX (residual jacobian, (N,N) numpy.ndarray). If config[‘reform_freq’] > 1, then fun should take two arguments The first is X, the second is a bool calc_grad where if True, fun returns a tuple of (R,dRdX). If False, fun just returns a tuple (R,) Function may return additional values in either tuple, but the additional values will be ignored here.

X0(N,) numpy.ndarray

Initial guess of the solution to the nonlinear problem.

verbosebool or None, optional

Flag to print convergence information if True. If None, then the configuration setting for ‘verbose’ from initialization will be used (default True) The default is None.

xtolfloat, optional

Tolerance to check for convergence on the step size. If None, then self.config[‘xtol’] is used. If that is also None, then 1e-6*X0.shape[0] is used as the x tolerance. Passing in a value here does not change the config value permanently (not parallel safe though). The default is None.

Returns:

X(N,) numpy.ndarray

Solution to the nonlinear problem that satisfies tolerances or from last step.

R(N,) numpy.ndarray

Residual vector from the last function evaluation (does not in generally correspond to value at X to save extra evaluation of fun).

dRdX(N,N) numpy.ndarray

Last residual Jacobian as evaluated during solution, not at final X.

soldict

Description of final convergence state. Has keys of [‘message’, ‘nfev’, ‘njev’, ‘success’]. ‘success’ is a bool with True corresponding to convergence. ‘nfev’ is the number of function evaluations completed. ‘njev’ is the number of jacobian evaluations. ‘message’ is either ‘Converged’ or ‘failed’. Use the bool from ‘success’ rather than the message for decisions. ‘nfev’ does not include any function evaluations done as part of the line search routine.

Other Parameters:

Dscalefloat or numpy.ndarray, optional

This argument does nothing. Conditioning of the numerical problem should be achieved by wrapping the problem of interest rather than here.

See also

NonlinearSolverOMP

Documentation for the solver class describes configurations and settings of the numerical solver that are configured at creation rather than at solution time.

line_search

Class method that may be used to improve convergence of nsolve if the appropriate solver settings are used.

Notes

This function uses either a full Newton-Raphson (NR) solver approach or Broyden-Fletcher-Goldfarb-Shanno (BFGS), which uses fewer NR iterations with some approximations of Jacobian between NR iterations. For BFGS see Algorithm 7.4 in [1].

The output status of current errors / norms of residual etc. are intended to provide easy checks to see what is going on during solutions. These outputs are slightly convoluted in that some outputs are for the previous iteration, check code for exact details.

References

[1]

Nocedal, J., S.J. Wright, 2006. Numerical optimization, 2nd ed, Springer series in operations research. Springer, New York.