jax_implicit_comp.py

An ImplicitComponent that uses JAX for derivatives.

class openmdao.components.jax_implicit_comp.JaxImplicitComponent(matrix_free=False, fallback_derivs_method='fd', **kwargs)

Bases: ImplicitComponent

Base class for implicit components when using JAX for derivatives.

Parameters:

matrix_freebool

If True, this component will compute derivatives using matrix vector products.

fallback_derivs_methodstr

The method to use if JAX is not available. Default is ‘fd’.

**kwargsdict

Additional arguments to be passed to the base class.

Attributes:

_tangentsdict

The tangents for the inputs and outputs.

_do_sparsitybool

If True, compute the sparsity.

_sparsitycoo_matrix or None

The sparsity of the Jacobian.

_jac_func_function or None

The function that computes the jacobian.

_jac_colored_function or None

The function that computes the colored jacobian.

_static_hashtuple

The hash of the static values.

_orig_compute_primalfunction

The original compute_primal method.

_ret_tuple_compute_primalfunction

The compute_primal method that returns a tuple.

_output_shapesdict

A dict of output shapes used when shapes are computed dynamically.

_do_shape_checkbool

If True, check the declared output shapes vs. the shapes of the outputs returned from compute_primal.

Methods

abs_meta_iter(iotype[, local, cont, discrete])	Iterate over absolute variable names and their metadata for this System.
add_constraint(name[, lower, upper, equals, ...])	Add a constraint variable to this system.
add_design_var(name[, lower, upper, ref, ...])	Add a design variable to this system.
add_discrete_input(name, val[, desc, tags, ...])	Add a discrete input variable to the component.
add_discrete_output(name, val[, desc, tags, ...])	Add an output variable to the component.
add_input(name, **kwargs)	Add an input to the component.
add_objective(name[, ref, ref0, index, ...])	Add a response variable to this system.
add_output(name, **kwargs)	Add an output to the component.
add_recorder(recorder[, recurse])	Add a recorder to the system.
add_residual(name[, shape, units, desc, ref])	Add a residual variable to the component.
add_response(name, type_[, lower, upper, ...])	Add a response variable to this system.
apply_linear(inputs, outputs, d_inputs, ...)	Compute jac-vector product.
apply_nonlinear(inputs, outputs, residuals)	Compute residuals given inputs and outputs.
best_partial_deriv_direction()	Return the best direction for partial deriv calculations based on input and output sizes.
check_config(logger)	Perform optional error checks.
check_partials([out_stream, compact_print, ...])	Check partial derivatives comprehensively for this component.
check_sparsity([method, max_nz, out_stream])	Check the sparsity of the computed jacobian against the declared sparsity.
cleanup()	Clean up resources prior to exit.
comm_info_iter()	Yield comm size for this system and all subsystems.
compute_fd_jac(jac[, method])	Force the use of finite difference to compute a jacobian.
compute_fd_sparsity([method, num_full_jacs, ...])	Use finite difference to compute a sparsity matrix.
compute_sparsity([direction, num_iters, ...])	Get the sparsity of the Jacobian.
convert2units(name, val, units)	Convert the given value to the specified units.
convert_from_units(name, val, units)	Convert the given value from the specified units to those of the named variable.
convert_units(name, val, units_from, units_to)	Wrap the utility convert_units and give a good error message.
declare_coloring(**kwargs)	Declare coloring for this component.
declare_partials(of, wrt[, dependent, rows, ...])	Declare information about this component's subjacobians.
dist_size_iter(io, top_comm)	Yield names and distributed ranges of all local and remote variables in this system.
get_coloring_fname(mode)	Return the full pathname to a coloring file.
get_conn_graph()	Return the model connection graph.
get_constraints([recurse, get_sizes, ...])	Get the Constraint settings from this system.
get_declare_partials_calls([sparsity])	Return a string containing declare_partials() calls based on the subjac sparsity.
get_design_vars([recurse, get_sizes, ...])	Get the DesignVariable settings from this system.
get_io_metadata([iotypes, metadata_keys, ...])	Retrieve metadata for a filtered list of variables.
get_linear_vectors()	Return the linear inputs, outputs, and residuals vectors.
get_nonlinear_vectors()	Return the inputs, outputs, and residuals vectors.
get_objectives([recurse, get_sizes, ...])	Get the Objective settings from this system.
get_outputs_dir(*subdirs[, mkdir])	Get the path under which all output files of this system are to be placed.
get_promotions([inprom, outprom])	Return all promotions for the given promoted variable(s).
get_reports_dir()	Get the path to the directory where the report files should go.
get_responses([recurse, get_sizes, use_prom_ivc])	Get the response variable settings from this system.
get_self_statics()	Override this in derived classes if compute_primal references static values.
get_source(name)	Return the source variable connected to the given named variable.
get_val(name[, units, indices, get_remote, ...])	Get an output/input/residual variable.
get_var_dup_info(name, io)	Return information about how the given variable is duplicated across MPI processes.
get_var_sizes(name, io)	Return the sizes of the given variable on all procs.
guess_nonlinear(inputs, outputs, residuals)	Provide initial guess for states.
has_vectors()	Check if the system vectors have been initialized.
initialize()	Perform any one-time initialization run at instantiation.
is_explicit([is_comp])	Return True if this is an explicit component.
linearize(inputs, outputs, jacobian[, ...])	Compute sub-jacobian parts and any applicable matrix factorizations.
list_inputs([val, prom_name, units, shape, ...])	Write a list of input names and other optional information to a specified stream.
list_options([include_default, ...])	Write a list of output names and other optional information to a specified stream.
list_outputs([explicit, implicit, val, ...])	Write a list of output names and other optional information to a specified stream.
list_vars([val, prom_name, residuals, ...])	Write a list of inputs and outputs sorted by component in execution order.
load_case(case)	Pull all input and output variables from a Case into this System.
load_model_options()	Load the relevant model options from Problem._metadata['model_options'].
override_method(name, method)	Dynamically add a method to this component instance.
record_iteration()	Record an iteration of the current System.
run_apply_linear(mode[, scope_out, scope_in])	Compute jac-vec product.
run_apply_nonlinear()	Compute residuals.
run_linearize([sub_do_ln])	Compute jacobian / factorization.
run_solve_linear(mode)	Apply inverse jac product.
run_solve_nonlinear()	Compute outputs.
run_validation()	Run validate method on all systems below this system.
set_check_partial_options(wrt[, method, ...])	Set options that will be used for checking partial derivatives.
set_constraint_options(name[, ref, ref0, ...])	Set options for constraints in the model.
set_design_var_options(name[, lower, upper, ...])	Set options for design vars in the model.
set_objective_options(name[, ref, ref0, ...])	Set options for objectives in the model.
set_output_solver_options(name[, lower, ...])	Set solver output options.
set_solver_print([level, depth, type_, ...])	Control printing for solvers and subsolvers in the model.
set_val(name, val[, units, indices])	Set an input or output variable.
setup()	Declare inputs and outputs.
setup_partials()	Declare partials.
setup_residuals()	User hook for adding named residuals to this component.
solve_linear(d_outputs, d_residuals, mode)	Apply inverse jac product.
solve_nonlinear(inputs, outputs)	Compute outputs given inputs.
sparsity_matches_fd([direction, outstream])	Compare the sparsity computed by this system vs.
subjac_sparsity_iter(sparsity[, wrt_matches])	Iterate over sparsity for each subjac in the jacobian.
system_iter([include_self, recurse, typ, ...])	Yield a generator of local subsystems of this system.
total_local_size(io)	Return the total local size of the given variable.
use_fixed_coloring([coloring, recurse])	Use a precomputed coloring for this System.
uses_approx()	Return True if the system uses approximations to compute derivatives.
validate(inputs, outputs[, discrete_inputs, ...])	Check any final input / output values after a run.

__init__(matrix_free=False, fallback_derivs_method='fd', **kwargs)

Store some bound methods so we can detect runtime overrides.

add_input(name, **kwargs)

Add an input to the component.

This overrides the base class method to update the kwargs to use dynamic shaping by default.

Parameters:

namestr

The name of the input.

**kwargsdict

The kwargs to pass to the base class method.

add_output(name, **kwargs)

Add an output to the component.

This overrides the base class method to update the kwargs to use dynamic shaping by default.

Parameters:

namestr

The name of the output.

**kwargsdict

The kwargs to pass to the base class method.

compute_sparsity(direction=None, num_iters=1, perturb_size=1e-09)

Get the sparsity of the Jacobian.

Parameters:

directionstr

The direction to compute the sparsity for.

num_itersint

The number of times to run the perturbation iteration.

perturb_sizefloat

The size of the perturbation to use.

Returns:

coo_matrix

The sparsity of the Jacobian.

declare_coloring(**kwargs)

Declare coloring for this component.

The ‘method’ argument is set to ‘jax’ and passed to the base class.

Parameters:

**kwargsdict

Additional arguments to be passed to the base class.