quimb.tensor.optimize ===================== .. py:module:: quimb.tensor.optimize .. autoapi-nested-parse:: Support for optimizing tensor networks using automatic differentiation to automatically derive gradients for input to scipy optimizers. Attributes ---------- .. autoapisummary:: quimb.tensor.optimize._DEFAULT_BACKEND quimb.tensor.optimize._REAL_CONVERSION quimb.tensor.optimize._COMPLEX_CONVERSION quimb.tensor.optimize._VARIABLE_TAG quimb.tensor.optimize.variable_finder quimb.tensor.optimize._BACKEND_HANDLERS quimb.tensor.optimize._STOC_GRAD_METHODS Classes ------- .. autoapisummary:: quimb.tensor.optimize.ArrayInfo quimb.tensor.optimize.Vectorizer quimb.tensor.optimize.AutoGradHandler quimb.tensor.optimize.JaxHandler quimb.tensor.optimize.TensorFlowHandler quimb.tensor.optimize.TorchHandler quimb.tensor.optimize.MultiLossHandler quimb.tensor.optimize.SGD quimb.tensor.optimize.RMSPROP quimb.tensor.optimize.ADAM quimb.tensor.optimize.NADAM quimb.tensor.optimize.ADABELIEF quimb.tensor.optimize.MakeArrayFn quimb.tensor.optimize.TNOptimizer Functions --------- .. autoapisummary:: quimb.tensor.optimize._parse_opt_in quimb.tensor.optimize._parse_opt_out quimb.tensor.optimize._parse_pytree_to_backend quimb.tensor.optimize.parse_network_to_backend quimb.tensor.optimize._inject_variables_pytree quimb.tensor.optimize.inject_variables quimb.tensor.optimize.convert_raw_arrays quimb.tensor.optimize.convert_variables_to_numpy quimb.tensor.optimize.get_autograd quimb.tensor.optimize.get_tensorflow quimb.tensor.optimize.get_torch quimb.tensor.optimize.identity_fn Module Contents --------------- .. py:data:: _DEFAULT_BACKEND :value: 'jax' .. py:data:: _REAL_CONVERSION .. py:data:: _COMPLEX_CONVERSION .. py:class:: ArrayInfo(array) Simple container for recording size and dtype information about arrays. .. py:attribute:: __slots__ :value: ('shape', 'size', 'dtype', 'iscomplex', 'real_size', 'equivalent_real_type', 'equivalent_complex_type') .. py:attribute:: shape .. py:attribute:: size .. py:attribute:: dtype .. py:attribute:: equivalent_real_type .. py:attribute:: equivalent_complex_type .. py:attribute:: iscomplex .. py:attribute:: real_size .. py:method:: __repr__() .. py:class:: Vectorizer(tree) Object for mapping back and forth between any pytree of mixed real/complex n-dimensional arrays to a single, real, double precision numpy vector, as required by ``scipy.optimize`` routines. :param tree: Any nested container of arrays, which will be flattened and packed into a single float64 vector. :type tree: pytree of array :param is_leaf: A function which takes a single argument and returns ``True`` if it is a leaf node in the tree and should be extracted, ``False`` otherwise. Defaults to everything that is not a tuple, list or dict. :type is_leaf: callable, optional .. py:attribute:: infos :value: [] .. py:attribute:: d :value: 0 .. py:attribute:: ref_tree .. py:method:: pack(tree, name='vector') Take ``arrays`` and pack their values into attribute `.{name}`, by default `.vector`. .. py:method:: unpack(vector=None) Turn the single, flat ``vector`` into a sequence of arrays. .. py:data:: _VARIABLE_TAG :value: '__VARIABLE{}__' .. py:data:: variable_finder .. py:function:: _parse_opt_in(tn, tags, shared_tags, to_constant) Parse a tensor network where tensors are assumed to be constant unless tagged. .. py:function:: _parse_opt_out(tn, constant_tags, to_constant) Parse a tensor network where tensors are assumed to be variables unless tagged. .. py:function:: _parse_pytree_to_backend(x, to_constant) Parse a arbitrary pytree, collecting variables. There is not opting in or out, all networks, tensors and raw arrays are considered variables. .. py:function:: parse_network_to_backend(tn, to_constant, tags=None, shared_tags=None, constant_tags=None) Parse tensor network to: - identify the dimension of the optimisation space and the initial point of the optimisation from the current values in the tensor network, - add variable tags to individual tensors so that optimisation vector values can be efficiently reinserted into the tensor network. There are two different modes: - 'opt in' : `tags` (and optionally `shared_tags`) are specified and only these tensor tags will be optimised over. In this case `constant_tags` is ignored if it is passed, - 'opt out' : `tags` is not specified. In this case all tensors will be optimised over, unless they have one of `constant_tags` tags. :param tn: The initial tensor network to parse. :type tn: TensorNetwork :param to_constant: Function that fixes a tensor as constant. :type to_constant: Callable :param tags: Set of opt-in tags to optimise. :type tags: str, or sequence of str, optional :param shared_tags: Subset of opt-in tags to joint optimise i.e. all tensors with tag s in shared_tags will correspond to the same optimisation variables. :type shared_tags: str, or sequence of str, optional :param constant_tags: Set of opt-out tags if `tags` not passed. :type constant_tags: str, or sequence of str, optional :returns: * **tn_ag** (*TensorNetwork*) -- Tensor network tagged for reinsertion of optimisation variable values. * **variables** (*list*) -- List of variables extracted from ``tn``. .. py:function:: _inject_variables_pytree(arrays, tree) .. py:function:: inject_variables(arrays, tn) Given the list of optimized variables ``arrays`` and the target tensor network or pytree ``tn``, inject the variables back in. .. py:function:: convert_raw_arrays(x, f) Given a ``TensorNetwork``, ``Tensor``, or other possibly structured raw array, return a copy where the underyling data has had ``f`` applied to it. Structured raw arrays should implement the ``tree = get_params()`` and ``set_params(tree)`` methods which get or set their underlying data using an arbitrary pytree. .. py:function:: convert_variables_to_numpy(x) .. py:function:: get_autograd() .. py:class:: AutoGradHandler(device='cpu') .. py:method:: to_variable(x) .. py:method:: to_constant(x) .. py:method:: setup_fn(fn) .. py:method:: value(arrays) .. py:method:: value_and_grad(arrays) .. py:class:: JaxHandler(jit_fn=True, device=None) .. py:attribute:: jit_fn :value: True .. py:attribute:: device :value: None .. py:method:: to_variable(x) .. py:method:: to_constant(x) .. py:method:: setup_fn(fn) .. py:method:: _setup_hessp(fn) .. py:method:: value(arrays) .. py:method:: value_and_grad(arrays) .. py:method:: hessp(primals, tangents) .. py:function:: get_tensorflow() .. py:class:: TensorFlowHandler(jit_fn=False, autograph=False, experimental_compile=False, device=None) .. py:attribute:: jit_fn :value: False .. py:attribute:: autograph :value: False .. py:attribute:: experimental_compile :value: False .. py:attribute:: device :value: None .. py:method:: to_variable(x) .. py:method:: to_constant(x) .. py:method:: setup_fn(fn) .. py:method:: value(arrays) .. py:method:: value_and_grad(arrays) .. py:function:: get_torch() .. py:class:: TorchHandler(jit_fn=False, device=None) .. py:attribute:: jit_fn :value: False .. py:attribute:: device :value: None .. py:method:: to_variable(x) .. py:method:: to_constant(x) .. py:method:: setup_fn(fn) .. py:method:: _setup_backend_fn(arrays) .. py:method:: value(arrays) .. py:method:: value_and_grad(arrays) .. py:data:: _BACKEND_HANDLERS .. py:class:: MultiLossHandler(autodiff_backend, executor=None, **backend_opts) .. py:attribute:: autodiff_backend .. py:attribute:: backend_opts .. py:attribute:: executor :value: None .. py:attribute:: handlers .. py:attribute:: to_constant .. py:method:: setup_fn(funcs) .. py:method:: _value_seq(arrays) .. py:method:: _value_par_seq(arrays) .. py:method:: value(arrays) .. py:method:: _value_and_grad_seq(arrays) .. py:method:: _value_and_grad_par(arrays) .. py:method:: value_and_grad(arrays) .. py:class:: SGD Stateful ``scipy.optimize.minimize`` compatible implementation of stochastic gradient descent with momentum. Adapted from ``autograd/misc/optimizers.py``. .. py:attribute:: OptimizeResult .. py:attribute:: _i :value: 0 .. py:attribute:: _velocity :value: None .. py:method:: get_velocity(x) .. py:method:: __call__(fun, x0, jac, args=(), learning_rate=0.1, mass=0.9, maxiter=1000, callback=None, bounds=None, **kwargs) .. py:class:: RMSPROP Stateful ``scipy.optimize.minimize`` compatible implementation of root mean squared prop: See Adagrad paper for details. Adapted from ``autograd/misc/optimizers.py``. .. py:attribute:: OptimizeResult .. py:attribute:: _i :value: 0 .. py:attribute:: _avg_sq_grad :value: None .. py:method:: get_avg_sq_grad(x) .. py:method:: __call__(fun, x0, jac, args=(), learning_rate=0.1, gamma=0.9, eps=1e-08, maxiter=1000, callback=None, bounds=None, **kwargs) .. py:class:: ADAM Stateful ``scipy.optimize.minimize`` compatible implementation of ADAM - http://arxiv.org/pdf/1412.6980.pdf. Adapted from ``autograd/misc/optimizers.py``. .. py:attribute:: OptimizeResult .. py:attribute:: _i :value: 0 .. py:attribute:: _m :value: None .. py:attribute:: _v :value: None .. py:method:: get_m(x) .. py:method:: get_v(x) .. py:method:: __call__(fun, x0, jac, args=(), learning_rate=0.001, beta1=0.9, beta2=0.999, eps=1e-08, maxiter=1000, callback=None, bounds=None, **kwargs) .. py:class:: NADAM Stateful ``scipy.optimize.minimize`` compatible implementation of NADAM - [Dozat - http://cs229.stanford.edu/proj2015/054_report.pdf]. Adapted from ``autograd/misc/optimizers.py``. .. py:attribute:: OptimizeResult .. py:attribute:: _i :value: 0 .. py:attribute:: _m :value: None .. py:attribute:: _v :value: None .. py:attribute:: _mus :value: None .. py:method:: get_m(x) .. py:method:: get_v(x) .. py:method:: get_mus(beta1) .. py:method:: __call__(fun, x0, jac, args=(), learning_rate=0.001, beta1=0.9, beta2=0.999, eps=1e-08, maxiter=1000, callback=None, bounds=None, **kwargs) .. py:class:: ADABELIEF Stateful ``scipy.optimize.minimize`` compatible implementation of ADABELIEF - https://arxiv.org/abs/2010.07468. Adapted from ``autograd/misc/optimizers.py``. .. py:attribute:: OptimizeResult .. py:attribute:: _i :value: 0 .. py:attribute:: _m :value: None .. py:attribute:: _s :value: None .. py:method:: get_m(x) .. py:method:: get_s(x) .. py:method:: __call__(fun, x0, jac, args=(), learning_rate=0.001, beta1=0.9, beta2=0.999, eps=1e-08, maxiter=1000, callback=None, bounds=None, **kwargs) .. py:data:: _STOC_GRAD_METHODS .. py:class:: MakeArrayFn(tn_opt, loss_fn, norm_fn, autodiff_backend) Class wrapper so picklable. .. py:attribute:: __name__ :value: 'MakeArrayFn' .. py:attribute:: tn_opt .. py:attribute:: loss_fn .. py:attribute:: norm_fn .. py:attribute:: autodiff_backend .. py:method:: __call__(arrays) .. py:function:: identity_fn(x) .. py:class:: TNOptimizer(tn, loss_fn, norm_fn=None, loss_constants=None, loss_kwargs=None, tags=None, shared_tags=None, constant_tags=None, loss_target=None, optimizer='L-BFGS-B', progbar=True, bounds=None, autodiff_backend='AUTO', executor=None, callback=None, **backend_opts) Globally optimize tensors within a tensor network with respect to any loss function via automatic differentiation. If parametrized tensors are used, optimize the parameters rather than the raw arrays. :param tn: The core tensor network structure within which to optimize tensors. :type tn: TensorNetwork :param loss_fn: The function that takes ``tn`` (as well as ``loss_constants`` and ``loss_kwargs``) and returns a single real 'loss' to be minimized. For Hamiltonians which can be represented as a sum over terms, an iterable collection of terms (e.g. list) can be given instead. In that case each term is evaluated independently and the sum taken as loss_fn. This can reduce the total memory requirements or allow for parallelization (see ``executor``). :type loss_fn: callable or sequence of callable :param norm_fn: A function to call before ``loss_fn`` that prepares or 'normalizes' the raw tensor network in some way. :type norm_fn: callable, optional :param loss_constants: Extra tensor networks, tensors, dicts/list/tuples of arrays, or arrays which will be supplied to ``loss_fn`` but also converted to the correct backend array type. :type loss_constants: dict, optional :param loss_kwargs: Extra options to supply to ``loss_fn`` (unlike ``loss_constants`` these are assumed to be simple options that don't need conversion). :type loss_kwargs: dict, optional :param tags: If supplied, only optimize tensors with any of these tags. :type tags: str, or sequence of str, optional :param shared_tags: If supplied, each tag in ``shared_tags`` corresponds to a group of tensors to be optimized together. :type shared_tags: str, or sequence of str, optional :param constant_tags: If supplied, skip optimizing tensors with any of these tags. This 'opt-out' mode is overridden if either ``tags`` or ``shared_tags`` is supplied. :type constant_tags: str, or sequence of str, optional :param loss_target: Stop optimizing once this loss value is reached. :type loss_target: float, optional :param optimizer: Which ``scipy.optimize.minimize`` optimizer to use (the ``'method'`` kwarg of that function). In addition, ``quimb`` implements a few custom optimizers compatible with this interface that you can reference by name - ``{'adam', 'nadam', 'rmsprop', 'sgd'}``. :type optimizer: str, optional :param executor: To be used with term-by-term Hamiltonians. If supplied, this executor is used to parallelize the evaluation. Otherwise each term is evaluated in sequence. It should implement the basic concurrent.futures (PEP 3148) interface. :type executor: None or Executor, optional :param progbar: Whether to show live progress. :type progbar: bool, optional :param bounds: Constrain the optimized tensor entries within this range (if the scipy optimizer supports it). :type bounds: None or (float, float), optional :param autodiff_backend: Which backend library to use to perform the automatic differentation (and computation). :type autodiff_backend: {'jax', 'autograd', 'tensorflow', 'torch'}, optional :param callback: A function to call after each optimization step. It should take the current ``TNOptimizer`` instance as its only argument. Information such as the current loss and number of evaluations can then be accessed:: def callback(tnopt): print(tnopt.nevals, tnopt.loss) :type callback: callable, optional :param backend_opts: Supplied to the backend function compiler and array handler. For example ``jit_fn=True`` or ``device='cpu'`` . .. py:attribute:: progbar :value: True .. py:attribute:: tags :value: None .. py:attribute:: shared_tags :value: None .. py:attribute:: constant_tags :value: None .. py:attribute:: _autodiff_backend :value: 'AUTO' .. py:attribute:: _multiloss .. py:attribute:: norm_fn :value: None .. py:attribute:: loss_constants .. py:attribute:: loss_kwargs .. py:property:: bounds .. py:property:: optimizer The underlying optimizer that works with the vectorized functions. .. py:attribute:: callback :value: None .. py:method:: _set_tn(tn) .. py:method:: _reset_tracking_info(loss_target=None) .. py:method:: reset(tn=None, clear_info=True, loss_target=None) Reset this optimizer without losing the compiled loss and gradient functions. :param tn: Set this tensor network as the current state of the optimizer, it must exactly match the original tensor network. :type tn: TensorNetwork, optional :param clear_info: Clear the tracked losses and iterations. :type clear_info: bool, optional .. py:method:: _maybe_init_pbar(n) .. py:method:: _maybe_update_pbar() .. py:method:: _maybe_close_pbar() .. py:method:: _check_loss_target() .. py:method:: _maybe_call_callback() .. py:method:: vectorized_value(x) The value of the loss function at vector ``x``. .. py:method:: vectorized_value_and_grad(x) The value and gradient of the loss function at vector ``x``. .. py:method:: vectorized_hessp(x, p) The action of the hessian at point ``x`` on vector ``p``. .. py:method:: __repr__() .. py:property:: d .. py:property:: nevals The number of gradient evaluations. .. py:method:: get_tn_opt() Extract the optimized tensor network, this is a three part process: 1. inject the current optimized vector into the target tensor network or pytree, 2. run it through ``norm_fn``, 3. drop any tags used to identify variables. :returns: **tn_opt** :rtype: TensorNetwork .. py:method:: optimize(n, tol=None, jac=True, hessp=False, optlib='scipy', **options) Run the optimizer for ``n`` function evaluations, using by default :func:`scipy.optimize.minimize` as the driver for the vectorized computation. Supplying the gradient and hessian vector product is controlled by the ``jac`` and ``hessp`` options respectively. :param n: Notionally the maximum number of iterations for the optimizer, note that depending on the optimizer being used, this may correspond to number of function evaluations rather than just iterations. :type n: int :param tol: Tolerance for convergence, note that various more specific tolerances can usually be supplied to ``options``, depending on the optimizer being used. :type tol: None or float, optional :param jac: Whether to supply the jacobian, i.e. gradient, of the loss function. :type jac: bool, optional :param hessp: Whether to supply the hessian vector product of the loss function. :type hessp: bool, optional :param optlib: Which optimization library to use. :type optlib: {'scipy', 'nlopt'}, optional :param options: Supplied to :func:`scipy.optimize.minimize` or whichever optimizer is being used. :returns: **tn_opt** :rtype: TensorNetwork .. py:method:: optimize_scipy(n, tol=None, jac=True, hessp=False, **options) Scipy based optimization, see :meth:`~quimb.tensor.optimize.TNOptimizer.optimize` for details. .. py:method:: optimize_basinhopping(n, nhop, temperature=1.0, jac=True, hessp=False, **options) Run the optimizer for using :func:`scipy.optimize.basinhopping` as the driver for the vectorized computation. This performs ``nhop`` local optimization each with ``n`` iterations. :param n: Number of iterations per local optimization. :type n: int :param nhop: Number of local optimizations to hop between. :type nhop: int :param temperature: H :type temperature: float, optional :param options: Supplied to the inner :func:`scipy.optimize.minimize` call. :returns: **tn_opt** :rtype: TensorNetwork .. py:method:: optimize_nlopt(n, tol=None, jac=True, hessp=False, ftol_rel=None, ftol_abs=None, xtol_rel=None, xtol_abs=None) Run the optimizer for ``n`` function evaluations, using ``nlopt`` as the backend library to run the optimization. Whether the gradient is computed depends on which ``optimizer`` is selected, see valid options at https://nlopt.readthedocs.io/en/latest/NLopt_Algorithms/. The following scipy ``optimizer`` options are automatically translated to the corresponding ``nlopt`` algorithms: {"l-bfgs-b", "slsqp", "tnc", "cobyla"}. :param n: The maximum number of iterations for the optimizer. :type n: int :param tol: Tolerance for convergence, here this is taken to be the relative tolerance for the loss (``ftol_rel`` below overrides this). :type tol: None or float, optional :param jac: Whether to supply the jacobian, i.e. gradient, of the loss function. :type jac: bool, optional :param hessp: Whether to supply the hessian vector product of the loss function. :type hessp: bool, optional :param ftol_rel: Set relative tolerance on function value. :type ftol_rel: float, optional :param ftol_abs: Set absolute tolerance on function value. :type ftol_abs: float, optional :param xtol_rel: Set relative tolerance on optimization parameters. :type xtol_rel: float, optional :param xtol_abs: Set absolute tolerances on optimization parameters. :type xtol_abs: float, optional :returns: **tn_opt** :rtype: TensorNetwork .. py:method:: optimize_ipopt(n, tol=None, **options) Run the optimizer for ``n`` function evaluations, using ``ipopt`` as the backend library to run the optimization via the python package ``cyipopt``. :param n: The maximum number of iterations for the optimizer. :type n: int :returns: **tn_opt** :rtype: TensorNetwork .. py:method:: optimize_nevergrad(n) Run the optimizer for ``n`` function evaluations, using ``nevergrad`` as the backend library to run the optimization. As the name suggests, the gradient is not required for this method. :param n: The maximum number of iterations for the optimizer. :type n: int :returns: **tn_opt** :rtype: TensorNetwork .. py:method:: plot(xscale='symlog', xscale_linthresh=20, zoom='auto', hlines=()) Plot the loss function as a function of the number of iterations. :param xscale: The scale of the x-axis. Default is ``"symlog"``, i.e. linear for the first part of the plot, and logarithmic for the rest, changing at ``xscale_linthresh``. :type xscale: str, optional :param xscale_linthresh: The threshold for the change from linear to logarithmic scale, if ``xscale`` is ``"symlog"``. Default is ``20``. :type xscale_linthresh: int, optional :param zoom: If not ``None``, show an inset plot of the last ``zoom`` iterations. :type zoom: None or int, optional :param hlines: A dictionary of horizontal lines to plot. The keys are the labels of the lines, and the values are the y-values of the lines. :type hlines: dict, optional :returns: * **fig** (*matplotlib.figure.Figure*) -- The figure object. * **ax** (*matplotlib.axes.Axes*) -- The axes object.