jax_sn.tools.parametrisation_tools module

class ParametrisationValues

Bases: Module

obj_constraint(theta, mod)

class ParametrisedObjective

Bases: Module

abstractmethod objective_and_constraints(theta: ParametrisationValues)

Return type:

Tuple[Union[float, PyTree]]

grad_objective_and_constraints(theta: ParametrisationValues)

Get the gradient of the objective and constraints with respect to the parametrisation values.

Uses JAX reverse mode autodiff.

Parameters:

theta: The parametrisation values.

Returns:

:

The gradient of the objective and constraints.

Return type:

Tuple[Union[float, PyTree]]

class ParametrisedModuleTheta(base_module: Module)

Bases: Module

Abstract base class for module problem parametrisations.

A problem parametrisation defines how to go from a set of parametrisation values (theta) and a base module (containing e.g. settings and non-parametrised data) to a parametrised module.

This adds a layer of abstraction on top of the module parametrisation. This is useful for example when the parametrisation values do not directly correspond to the module parametrisation theta, but need to be transformed first (e.g. when the parametrisation values are in log-space, or when the parametrisation values correspond to a different representation of the same data).

A concrete example is that for a jax_sn.operators.multi_group_operator.MultiGroupOperator, the module parametrisation may be given by MultigroupOperatorParametrisedVertices, (i.e. in terms of the vertices) but the problem parametrisation may define a mapping from a set of optimisation parameters to those vertices.

This provides a interface to go from the optimisation parameters to the module theta via the module parametrisation.

base_module: Module

The base module containing non-parametrised data and settings (and initial data possibly)

abstractmethod to_module(theta: ParametrisationValues)

Return type:

Module

class ParametrisedModuleObjective(theta_parametrised_module: ParametrisedModuleTheta)

Bases: ParametrisedObjective

Objective defined on a parametrised module via a parametrised module theta. This is useful to recycle functions that will be the same for different specific parametrisations but similar values.

A concrete example is that for a jax_sn.operators.multi_group_operator.MultiGroupOperator, although the parametrisation is different (handled by different ParametrisedModuleTheta subclasses) and the objective may be different (e.g. different ways to go from solution to loss), the standard linear solve from a non-parametrised right-hand-side is always the same.

theta_parametrised_module: ParametrisedModuleTheta

The parametrised module defining the mapping from parametrisation values to the module.

abstractmethod objective_and_constraints(theta: ParametrisationValues)

Return type:

Tuple[Union[float, PyTree]]

grad_objective_and_constraints(theta: ParametrisationValues)

Get the gradient of the objective and constraints with respect to the parametrisation values.

Uses JAX reverse mode autodiff.

Parameters:

theta: The parametrisation values.

Returns:

:

The gradient of the objective and constraints.

Return type:

Tuple[Union[float, PyTree]]

to_module(theta: ParametrisationValues)

Get the parametrised module corresponding to the given parametrisation values.

Parameters:

theta: The parametrisation values.

Returns:

:

The parametrised module.

Return type:

Module

class ParametrisedScipyInterface(parametrised_objective: ParametrisedObjective, f_theta_parametrisation_values: Tuple[Callable[[Array], ParametrisationValues], Callable[[ParametrisationValues], Array]])

Bases: object

parametrised_objective: ParametrisedObjective

f_theta_parametrisation_values: Callable[[Array], ParametrisationValues]

evaluate_values(theta)

evaluate_jac(theta)

objective(x)

constraints(x)

grad_objective(x)

jac_constraints(x)

class DiagParametrisedScipyInterface(parametrised_objective: ParametrisedObjective, f_theta_parametrisation_values: Callable[[Array], ParametrisationValues], _cache_theta_val: Array = None, _cache_val: Array = None, _cache_theta_jac: Array = None, _cache_jac: Array = None, iteration_i: int = 0, n_in_iteration: int = 0, save_dir: str = 'iterations/')

Bases: ParametrisedScipyInterface

iteration_i: int = 0

n_in_iteration: int = 0

save_dir: str = 'iterations/'

classmethod from_scipy_interface(scipy_interface: ParametrisedScipyInterface)

evaluate_values(theta)

evaluate_jac(theta)

ks_agg(x, alpha=50.0)