def minimize(fcn: Callable[..., torch.Tensor], y0: torch.Tensor, params: Sequence[Any] = [], bck_options: Mapping[str, Any] = {}, method: Union[str, None] = None, **fwd_options) -> torch.Tensor: """ Solve the unbounded minimization problem: .. math:: \mathbf{y^*} = \\arg\min_\mathbf{y} f(\mathbf{y}, \\theta) to find the best :math:`\mathbf{y}` that minimizes the output of the function :math:`f`. Arguments --------- fcn: callable The function to be optimized with output tensor with 1 element. y0: torch.tensor Initial guess of the solution with shape ``(*ny)`` params: list List of any other parameters to be put in ``fcn`` bck_options: dict Method-specific options for the backward solve. method: str or None Minimization method. **fwd_options Method-specific options (see method section) Returns ------- torch.tensor The solution of the minimization with shape ``(*ny)`` """ # perform implementation check if debug mode is enabled if is_debug_enabled(): assert_fcn_params(fcn, (y0, *params)) pfunc = get_pure_function(fcn) fwd_options["method"] = _get_minimizer_default_method(method) # the rootfinder algorithms are designed to move to the opposite direction # of the output of the function, so the output of this function is just # the grad of z w.r.t. y @make_sibling(pfunc) def new_fcn(y, *params): with torch.enable_grad(): y1 = y.clone().requires_grad_() z = pfunc(y1, *params) grady, = torch.autograd.grad(z, (y1, ), retain_graph=True, create_graph=torch.is_grad_enabled()) return grady return _RootFinder.apply(new_fcn, y0, fwd_options, bck_options, len(params), *params, *pfunc.objparams())
def equilibrium(fcn: Callable[..., torch.Tensor], y0: torch.Tensor, params: Sequence[Any] = [], bck_options: Mapping[str, Any] = {}, method: Union[str, None] = None, **fwd_options): """ Solving the equilibrium equation of a given function, .. math:: \mathbf{y} = \mathbf{f}(\mathbf{y}, \\theta) where :math:`\mathbf{f}` is a function that can be non-linear and produce output of the same shape of :math:`\mathbf{y}`, and :math:`\\theta` is other parameters required in the function. The output of this block is :math:`\mathbf{y}` that produces the same :math:`\mathbf{y}` as the output. Arguments --------- fcn : callable The function :math:`\mathbf{f}` with output tensor ``(*ny)`` y0 : torch.tensor Initial guess of the solution with shape ``(*ny)`` params : list List of any other parameters to be put in ``fcn`` bck_options : dict Method-specific options for the backward solve method : str or None Rootfinder method. **fwd_options Method-specific options (see method section) Returns ------- torch.tensor The solution which satisfies :math:`\mathbf{y} = \mathbf{f}(\mathbf{y},\\theta)` with shape ``(*ny)`` """ # perform implementation check if debug mode is enabled if is_debug_enabled(): assert_fcn_params(fcn, (y0, *params)) pfunc = get_pure_function(fcn) @make_sibling(pfunc) def new_fcn(y, *params): return y - pfunc(y, *params) fwd_options["method"] = _get_rootfinder_default_method(method) return _RootFinder.apply(new_fcn, y0, fwd_options, bck_options, len(params), *params, *pfunc.getobjparams())
def minimize(fcn: Callable[..., torch.Tensor], y0: torch.Tensor, params: Sequence[Any] = [], fwd_options: Mapping[str, Any] = {}, bck_options: Mapping[str, Any] = {}) -> torch.Tensor: """ Solve the minimization problem: z = (argmin_y) fcn(y, *params) to find the best `y` that minimizes the output of the function `fcn`. The output of `fcn` must be a single element tensor. Arguments --------- * fcn: callable with output tensor (numel=1) The function * y0: torch.tensor with shape (*ny) Initial guess of the solution * params: list List of any other parameters to be put in fcn * fwd_options: dict Options for the minimizer method * bck_options: dict Options for the backward solve method Returns ------- * y: torch.tensor with shape (*ny) The solution of the minimization """ # perform implementation check if debug mode is enabled if is_debug_enabled(): assert_fcn_params(fcn, (y0, *params)) pfunc = get_pure_function(fcn) # the rootfinder algorithms are designed to move to the opposite direction # of the output of the function, so the output of this function is just # the grad of z w.r.t. y @make_sibling(pfunc) def new_fcn(y, *params): with torch.enable_grad(): y1 = y.clone().requires_grad_() z = pfunc(y1, *params) grady, = torch.autograd.grad(z, (y1, ), retain_graph=True, create_graph=torch.is_grad_enabled()) return grady return _RootFinder.apply(new_fcn, y0, fwd_options, bck_options, len(params), *params, *pfunc.objparams())
def equilibrium(fcn: Callable[..., torch.Tensor], y0: torch.Tensor, params: Sequence[Any] = [], fwd_options: Mapping[str, Any] = {}, bck_options: Mapping[str, Any] = {}): """ Solving the equilibrium equation of a given function, y = fcn(y, *params) where `fcn` is a function that can be non-linear and produce output of shape `y`. The output of this block is `y` that produces the 0 as the output. Arguments --------- * fcn: callable with output tensor (*ny) The function * y0: torch.tensor with shape (*ny) Initial guess of the solution * params: list List of any other parameters to be put in fcn * fwd_options: dict Options for the rootfinder method * bck_options: dict Options for the backward solve method Returns ------- * yout: torch.tensor with shape (*ny) The solution which satisfies yout = fcn(yout) Note ---- * To obtain the correct gradient and higher order gradients, the fcn must be: - a torch.nn.Module with fcn.parameters() list the tensors that determine the output of the fcn. - a method in xt.EditableModule object with no out-of-scope parameters. - a function with no out-of-scope parameters. """ # perform implementation check if debug mode is enabled if is_debug_enabled(): assert_fcn_params(fcn, (y0, *params)) pfunc = get_pure_function(fcn) @make_sibling(pfunc) def new_fcn(y, *params): return y - pfunc(y, *params) return _RootFinder.apply(new_fcn, y0, fwd_options, bck_options, len(params), *params, *pfunc.getobjparams())
def _mcquad(ffcn, log_pfcn, x0, xsamples, wsamples, fparams, pparams, method, bck_options, **fwd_options): # this is mcquad with an additional xsamples argument, to prevent xsamples being set by users if is_debug_enabled(): assert_fcn_params(ffcn, (x0, *fparams)) assert_fcn_params(log_pfcn, (x0, *pparams)) # check if ffcn produces a list / tuple out = ffcn(x0, *fparams) is_tuple_out = isinstance(out, list) or isinstance(out, tuple) # get the pure functions pure_ffcn = get_pure_function(ffcn) pure_logpfcn = get_pure_function(log_pfcn) nfparams = len(fparams) npparams = len(pparams) fobjparams = pure_ffcn.objparams() pobjparams = pure_logpfcn.objparams() nf_objparams = len(fobjparams) if is_tuple_out: packer = TensorPacker(out) @make_sibling(pure_ffcn) def pure_ffcn2(x, *fparams): y = pure_ffcn(x, *fparams) return packer.flatten(y) res = _MCQuad.apply(pure_ffcn2, pure_logpfcn, x0, None, None, method, fwd_options, bck_options, nfparams, nf_objparams, npparams, *fparams, *fobjparams, *pparams, *pobjparams) return packer.pack(res) else: return _MCQuad.apply(pure_ffcn, pure_logpfcn, x0, None, None, method, fwd_options, bck_options, nfparams, nf_objparams, npparams, *fparams, *fobjparams, *pparams, *pobjparams)
def solve_ivp(fcn: Union[Callable[..., torch.Tensor], Callable[..., Sequence[torch.Tensor]]], ts: torch.Tensor, y0: torch.Tensor, params: Sequence[Any] = [], bck_options: Mapping[str, Any] = {}, method: Union[str, Callable, None] = None, **fwd_options) -> Union[torch.Tensor, Sequence[torch.Tensor]]: r""" Solve the initial value problem (IVP) or also commonly known as ordinary differential equations (ODE), where given the initial value :math:`\mathbf{y_0}`, it then solves .. math:: \mathbf{y}(t) = \mathbf{y_0} + \int_{t_0}^{t} \mathbf{f}(t', \mathbf{y}, \theta)\ \mathrm{d}t' Arguments --------- fcn: callable The function that represents dy/dt. The function takes an input of a single time ``t`` and tensor ``y`` with shape ``(*ny)`` and produce :math:`\mathrm{d}\mathbf{y}/\mathrm{d}t` with shape ``(*ny)``. The output of the function must be a tensor with shape ``(*ny)`` or a list of tensors. ts: torch.tensor The time points where the value of `y` will be returned. It must be monotonically increasing or decreasing. It is a tensor with shape ``(nt,)``. y0: torch.tensor The initial value of ``y``, i.e. ``y(t[0]) == y0``. It is a tensor with shape ``(*ny)`` or a list of tensors. params: list Sequence of other parameters required in the function. bck_options: dict Options for the backward solve_ivp method. If not specified, it will take the same options as fwd_options. method: str or callable or None Initial value problem solver. If None, it will choose ``"rk45"``. **fwd_options Method-specific option (see method section below). Returns ------- torch.tensor or a list of tensors The values of ``y`` for each time step in ``ts``. It is a tensor with shape ``(nt,*ny)`` or a list of tensors """ if is_debug_enabled(): assert_fcn_params(fcn, (ts[0], y0, *params)) assert_runtime(len(ts.shape) == 1, "Argument ts must be a 1D tensor") if method is None: # set the default method method = "rk45" fwd_options["method"] = method # run once to see if the outputs is a tuple or a single tensor is_y0_list = isinstance(y0, list) or isinstance(y0, tuple) dydt = fcn(ts[0], y0, *params) is_dydt_list = isinstance(dydt, list) or isinstance(dydt, tuple) if is_y0_list != is_dydt_list: raise RuntimeError( "The y0 and output of fcn must both be tuple or a tensor") pfcn = get_pure_function(fcn) if is_y0_list: nt = len(ts) roller = TensorPacker(y0) @make_sibling(pfcn) def pfcn2(t, ytensor, *params): ylist = roller.pack(ytensor) res_list = pfcn(t, ylist, *params) res = roller.flatten(res_list) return res y0 = roller.flatten(y0) res = _SolveIVP.apply(pfcn2, ts, fwd_options, bck_options, len(params), y0, *params, *pfcn.objparams()) return roller.pack(res) else: return _SolveIVP.apply(pfcn, ts, fwd_options, bck_options, len(params), y0, *params, *pfcn.objparams())
def equilibrium(fcn: Callable[..., torch.Tensor], y0: torch.Tensor, params: Sequence[Any] = [], bck_options: Mapping[str, Any] = {}, method: Union[str, Callable, None] = None, **fwd_options) -> torch.Tensor: r""" Solving the equilibrium equation of a given function, .. math:: \mathbf{y} = \mathbf{f}(\mathbf{y}, \theta) where :math:`\mathbf{f}` is a function that can be non-linear and produce output of the same shape of :math:`\mathbf{y}`, and :math:`\theta` is other parameters required in the function. The output of this block is :math:`\mathbf{y}` that produces the same :math:`\mathbf{y}` as the output. Arguments --------- fcn : callable The function :math:`\mathbf{f}` with output tensor ``(*ny)`` y0 : torch.tensor Initial guess of the solution with shape ``(*ny)`` params : list Sequence of any other parameters to be put in ``fcn`` bck_options : dict Method-specific options for the backward solve (see :func:`xitorch.linalg.solve`) method : str or None Rootfinder method. If None, it will choose ``"broyden1"``. **fwd_options Method-specific options (see method section) Returns ------- torch.tensor The solution which satisfies :math:`\mathbf{y} = \mathbf{f}(\mathbf{y},\theta)` with shape ``(*ny)`` Example ------- .. testsetup:: equil1 import torch from xitorch.optimize import equilibrium .. doctest:: equil1 >>> def func1(y, A): # example function ... return torch.tanh(A @ y + 0.1) + y / 2.0 >>> A = torch.tensor([[1.1, 0.4], [0.3, 0.8]]).requires_grad_() >>> y0 = torch.zeros((2,1)) # zeros as the initial guess >>> yequil = equilibrium(func1, y0, params=(A,)) >>> print(yequil) tensor([[ 0.2313], [-0.5957]], grad_fn=<_RootFinderBackward>) Note ---- * This is a direct implementation of finding the root of :math:`\mathbf{g}(\mathbf{y}, \theta) = \mathbf{y} - \mathbf{f}(\mathbf{y}, \theta)` """ # perform implementation check if debug mode is enabled if is_debug_enabled(): assert_fcn_params(fcn, (y0, *params)) pfunc = get_pure_function(fcn) @make_sibling(pfunc) def new_fcn(y, *params): return y - pfunc(y, *params) fwd_options["method"] = _get_rootfinder_default_method(method) return _RootFinder.apply(new_fcn, y0, fwd_options, bck_options, len(params), *params, *pfunc.objparams())
def minimize(fcn: Callable[..., torch.Tensor], y0: torch.Tensor, params: Sequence[Any] = [], bck_options: Mapping[str, Any] = {}, method: Union[str, Callable] = None, **fwd_options) -> torch.Tensor: r""" Solve the unbounded minimization problem: .. math:: \mathbf{y^*} = \arg\min_\mathbf{y} f(\mathbf{y}, \theta) to find the best :math:`\mathbf{y}` that minimizes the output of the function :math:`f`. Arguments --------- fcn: callable The function to be optimized with output tensor with 1 element. y0: torch.tensor Initial guess of the solution with shape ``(*ny)`` params: list Sequence of any other parameters to be put in ``fcn`` bck_options: dict Method-specific options for the backward solve (see :func:`xitorch.linalg.solve`) method: str or callable or None Minimization method. If None, it will choose ``"broyden1"``. **fwd_options Method-specific options (see method section) Returns ------- torch.tensor The solution of the minimization with shape ``(*ny)`` Example ------- .. testsetup:: root1 import torch from xitorch.optimize import minimize .. doctest:: root1 >>> def func1(y, A): # example function ... return torch.sum((A @ y)**2 + y / 2.0) >>> A = torch.tensor([[1.1, 0.4], [0.3, 0.8]]).requires_grad_() >>> y0 = torch.zeros((2,1)) # zeros as the initial guess >>> ymin = minimize(func1, y0, params=(A,)) >>> print(ymin) tensor([[-0.0519], [-0.2684]], grad_fn=<_RootFinderBackward>) """ # perform implementation check if debug mode is enabled if is_debug_enabled(): assert_fcn_params(fcn, (y0, *params)) pfunc = get_pure_function(fcn) fwd_options["method"] = _get_minimizer_default_method(method) # the rootfinder algorithms are designed to move to the opposite direction # of the output of the function, so the output of this function is just # the grad of z w.r.t. y @make_sibling(pfunc) def new_fcn(y, *params): with torch.enable_grad(): y1 = y.clone().requires_grad_() z = pfunc(y1, *params) grady, = torch.autograd.grad(z, (y1, ), retain_graph=True, create_graph=torch.is_grad_enabled()) return grady return _RootFinder.apply(new_fcn, y0, fwd_options, bck_options, len(params), *params, *pfunc.objparams())
def quad(fcn: Union[Callable[..., torch.Tensor], Callable[..., List[torch.Tensor]]], xl: Union[float, int, torch.Tensor], xu: Union[float, int, torch.Tensor], params: Sequence[Any] = [], fwd_options: Mapping[str, Any] = {}, bck_options: Mapping[str, Any] = {}): """ Calculate the quadrature of the function `fcn` from `x0` to `xf`: y = int_xl^xu fcn(x, *params) Arguments --------- * fcn: callable with output tensor with shape (*nout) or list of tensors The function to be integrated. * xl, xu: float, int, or 1-element torch.Tensor The lower and upper bound of the integration. * params: list List of any other parameters for the function `fcn`. * fwd_options: dict Options for the forward quadrature method. * bck_options: dict Options for the backward quadrature method. Returns ------- * y: torch.tensor with shape (*nout) or list of tensors The quadrature results. """ # perform implementation check if debug mode is enabled if is_debug_enabled(): assert_fcn_params(fcn, (xl, *params)) if isinstance(xl, torch.Tensor): assert_runtime(torch.numel(xl) == 1, "xl must be a 1-element tensors") if isinstance(xu, torch.Tensor): assert_runtime(torch.numel(xu) == 1, "xu must be a 1-element tensors") out = fcn(xl, *params) is_tuple_out = not isinstance(out, torch.Tensor) if not is_tuple_out: dtype = out.dtype device = out.device elif len(out) > 0: dtype = out[0].dtype device = out[0].device else: raise RuntimeError("The output of the fcn must be non-empty") pfunc = get_pure_function(fcn) nparams = len(params) if is_tuple_out: packer = TensorPacker(out) @make_sibling(pfunc) def pfunc2(x, *params): y = fcn(x, *params) return packer.flatten(y) res = _Quadrature.apply(pfunc2, xl, xu, fwd_options, bck_options, nparams, dtype, device, *params, *pfunc.objparams()) return packer.pack(res) else: return _Quadrature.apply(pfunc, xl, xu, fwd_options, bck_options, nparams, dtype, device, *params, *pfunc.objparams())
def solve_ivp(fcn: Callable[..., torch.Tensor], ts: torch.Tensor, y0: torch.Tensor, params: Sequence[Any] = [], fwd_options: Mapping[str, Any] = {}, bck_options: Mapping[str, Any] = {}) -> torch.Tensor: """ Solve the initial value problem (IVP) which given the initial value `y0`, the function is then solve y(t) = y0 + int_t0^t f(t', y, *params) dt' Arguments --------- * fcn: callable with output a tensor with shape (*ny) or a list of tensors The function that represents dy/dt. The function takes an input of a single time `t` and `y` with shape (*ny) and produce dydt with shape (*ny). * ts: torch.tensor with shape (nt,) The time points where the value of `y` is returned. It must be monotonically increasing or decreasing. * y0: torch.tensor with shape (*ny) or a list of tensors The initial value of y, i.e. y(t[0]) == y0 * params: list List of other parameters required in the function. * fwd_options: dict Options for the forward solve_ivp method. * bck_options: dict Options for the backward solve_ivp method. Returns ------- * yt: torch.tensor with shape (nt,*ny) or a list of tensors The values of `y` for each time step in `ts`. """ if is_debug_enabled(): assert_fcn_params(fcn, (ts[0], y0, *params)) assert_runtime(len(ts.shape) == 1, "Argument ts must be a 1D tensor") # run once to see if the outputs is a tuple or a single tensor is_y0_list = isinstance(y0, list) or isinstance(y0, tuple) dydt = fcn(ts[0], y0, *params) is_dydt_list = isinstance(dydt, list) or isinstance(dydt, tuple) if is_y0_list != is_dydt_list: raise RuntimeError( "The y0 and output of fcn must both be tuple or a tensor") pfcn = get_pure_function(fcn) if is_y0_list: nt = len(ts) roller = TensorPacker(y0) @make_sibling(pfcn) def pfcn2(t, ytensor, *params): ylist = roller.pack(ytensor) res_list = pfcn(t, ylist, *params) res = roller.flatten(res_list) return res y0 = roller.flatten(y0) res = _SolveIVP.apply(pfcn2, ts, fwd_options, bck_options, len(params), y0, *params, *pfcn.objparams()) return roller.pack(res) else: return _SolveIVP.apply(pfcn, ts, fwd_options, bck_options, len(params), y0, *params, *pfcn.objparams())
def quad(fcn: Union[Callable[..., torch.Tensor], Callable[..., Sequence[torch.Tensor]]], xl: Union[float, int, torch.Tensor], xu: Union[float, int, torch.Tensor], params: Sequence[Any] = [], bck_options: Mapping[str, Any] = {}, method: Union[str, Callable, None] = None, **fwd_options) -> Union[torch.Tensor, Sequence[torch.Tensor]]: r""" Calculate the quadrature: .. math:: y = \int_{x_l}^{x_u} f(x, \theta)\ \mathrm{d}x Arguments --------- fcn: callable The function to be integrated. Its output must be a tensor with shape ``(*nout)`` or list of tensors. xl: float, int or 1-element torch.Tensor The lower bound of the integration. xu: float, int or 1-element torch.Tensor The upper bound of the integration. params: list Sequence of any other parameters for the function ``fcn``. bck_options: dict Options for the backward quadrature method. method: str or callable or None Quadrature method. If None, it will choose ``"leggauss"``. **fwd_options Method-specific options (see method section). Returns ------- torch.tensor or a list of tensors The quadrature results with shape ``(*nout)`` or list of tensors. """ # perform implementation check if debug mode is enabled if is_debug_enabled(): assert_fcn_params(fcn, (xl, *params)) if isinstance(xl, torch.Tensor): assert_runtime(torch.numel(xl) == 1, "xl must be a 1-element tensors") if isinstance(xu, torch.Tensor): assert_runtime(torch.numel(xu) == 1, "xu must be a 1-element tensors") if method is None: method = "leggauss" fwd_options["method"] = method out = fcn(xl, *params) if isinstance(out, torch.Tensor): dtype = out.dtype device = out.device is_tuple_out = False elif len(out) > 0: dtype = out[0].dtype device = out[0].device is_tuple_out = True else: raise RuntimeError("The output of the fcn must be non-empty") pfunc = get_pure_function(fcn) nparams = len(params) if is_tuple_out: packer = TensorPacker(out) @make_sibling(pfunc) def pfunc2(x, *params): y = fcn(x, *params) return packer.flatten(y) res = _Quadrature.apply(pfunc2, xl, xu, fwd_options, bck_options, nparams, dtype, device, *params, *pfunc.objparams()) return packer.pack(res) else: return _Quadrature.apply(pfunc, xl, xu, fwd_options, bck_options, nparams, dtype, device, *params, *pfunc.objparams())