def forward(ctx, pfcn, ts, fwd_options, bck_options, nparams, y0,
*allparams):
config = fwd_options
ctx.bck_config = set_default_option(config, bck_options)
params = allparams[:nparams]
objparams = allparams[nparams:]
method = config.pop("method")
methods = {
"rk4": rk4_ivp,
"rk38": rk38_ivp,
"rk23": rk23_adaptive,
"rk45": rk45_adaptive,
}
solver = get_method("solve_ivp", methods, method)
yt = solver(pfcn, ts, y0, params, **config)
# save the parameters for backward
ctx.param_sep = TensorNonTensorSeparator(allparams, varonly=True)
tensor_params = ctx.param_sep.get_tensor_params()
ctx.save_for_backward(ts, y0, *tensor_params)
ctx.pfcn = pfcn
ctx.nparams = nparams
ctx.yt = yt
ctx.ts_requires_grad = ts.requires_grad
return yt
def forward(ctx, A, neig, mode, M, fwd_options, bck_options, na,
*amparams):
# A: LinearOperator (*BA, q, q)
# M: LinearOperator (*BM, q, q) or None
# separate the sets of parameters
params = amparams[:na]
mparams = amparams[na:]
config = set_default_option({}, fwd_options)
ctx.bck_config = set_default_option(
{
# "method": ???
},
bck_options)
method = config.pop("method")
with A.uselinopparams(*params), M.uselinopparams(
*mparams) if M is not None else dummy_context_manager():
methods = {
"davidson": davidson,
"custom_exacteig": custom_exacteig,
}
method_fcn = get_method("symeig", methods, method)
evals, evecs = method_fcn(A, neig, mode, M, **config)
# save for the backward
ctx.evals = evals # (*BAM, neig)
ctx.evecs = evecs # (*BAM, na, neig)
ctx.params = params
ctx.A = A
ctx.M = M
ctx.mparams = mparams
return evals, evecs
def forward(ctx, fcn, y0, options, bck_options, nparams, *allparams):
# set default options
config = options
ctx.bck_options = bck_options
params = allparams[:nparams]
objparams = allparams[nparams:]
with fcn.useobjparams(objparams):
method = config.pop("method")
methods = {
"broyden1": broyden1,
"broyden2": broyden2,
"linearmixing": linearmixing,
}
method_fcn = get_method("rootfinder", methods, method)
y = method_fcn(fcn, y0, params, **config)
ctx.fcn = fcn
# split tensors and non-tensors params
ctx.nparams = nparams
ctx.param_sep = TensorNonTensorSeparator(allparams)
tensor_params = ctx.param_sep.get_tensor_params()
ctx.save_for_backward(y, *tensor_params)
return y
def forward(ctx, fcn, y0, fwd_fcn, is_opt_method, options, bck_options,
nparams, *allparams):
# fcn: a function that returns what has to be 0 (will be used in the
# backward, not used in the forward). For minimization, it is
# the gradient
# fwd_fcn: a function that will be executed in the forward method
# (unused in the backward)
# This class is also used for minimization, where fcn and fwd_fcn might
# be slightly different
# set default options
config = options
ctx.bck_options = bck_options
params = allparams[:nparams]
objparams = allparams[nparams:]
with fwd_fcn.useobjparams(objparams):
method = config.pop("method")
methods = _RF_METHODS if not is_opt_method else _OPT_METHODS
name = "rootfinder" if not is_opt_method else "minimizer"
method_fcn = get_method(name, methods, method)
y = method_fcn(fwd_fcn, y0, params, **config)
ctx.fcn = fcn
ctx.is_opt_method = is_opt_method
# split tensors and non-tensors params
ctx.nparams = nparams
ctx.param_sep = TensorNonTensorSeparator(allparams)
tensor_params = ctx.param_sep.get_tensor_params()
ctx.save_for_backward(y, *tensor_params)
return y
def __init__(self,
x: torch.Tensor,
y: Optional[torch.Tensor] = None,
method: Union[str, Callable, None] = None,
**fwd_options):
if method is None:
method = "cspline"
methods = {
"cspline": CubicSpline1D,
}
method_cls = get_method("Interp1D", methods, method)
self.obj = method_cls(x, y, **fwd_options)
def forward(ctx, fcn, xl, xu, fwd_options, bck_options, nparams, dtype,
device, *all_params):
with fcn.disable_state_change():
config = fwd_options
ctx.bck_config = set_default_option(config, bck_options)
params = all_params[:nparams]
objparams = all_params[nparams:]
# convert to tensor
xl = torch.as_tensor(xl, dtype=dtype, device=device)
xu = torch.as_tensor(xu, dtype=dtype, device=device)
# apply transformation if the boundaries contain inf
if _isinf(xl) or _isinf(xu):
tfm = _TanInfTransform()
@make_sibling(fcn)
def fcn2(t, *params):
ys = fcn(tfm.forward(t), *params)
dxdt = tfm.dxdt(t)
return ys * dxdt
tl = tfm.x2t(xl)
tu = tfm.x2t(xu)
else:
fcn2 = fcn
tl = xl
tu = xu
method = config.pop("method")
methods = {"leggauss": leggauss}
method_fcn = get_method("quad", methods, method)
y = method_fcn(fcn2, tl, tu, params, **config)
# save the parameters for backward
ctx.param_sep = TensorNonTensorSeparator(all_params)
tensor_params = ctx.param_sep.get_tensor_params()
ctx.xltensor = isinstance(xl, torch.Tensor)
ctx.xutensor = isinstance(xu, torch.Tensor)
xlxu_tensor = ([xl] if ctx.xltensor else []) + \
([xu] if ctx.xutensor else [])
ctx.xlxu_nontensor = ([xl] if not ctx.xltensor else []) + \
([xu] if not ctx.xutensor else [])
ctx.save_for_backward(*xlxu_tensor, *tensor_params)
ctx.fcn = fcn
ctx.nparams = nparams
return y
def __init__(self, x: torch.Tensor,
method: Union[str, Callable, None] = None,
**fwd_options):
if method is None:
method = "cspline"
if not (isinstance(x, torch.Tensor) and len(x.shape) == 1):
raise RuntimeError("The input x to SQuad must be a 1D tensor")
all_clss = {
"cspline": CubicSplineSQuad,
"simpson": SimpsonSQuad,
"trapz": TrapzSQuad,
}
clss = get_method("SQuad", all_clss, method)
self.obj = clss(x, **fwd_options)
self.nx = x.shape[-1]
def forward(ctx, ffcn, log_pfcn, x0, xsamples, wsamples, method,
fwd_options, bck_options, nfparams, nf_objparams, npparams,
*all_fpparams):
# set up the default options
config = fwd_options
ctx.bck_config = set_default_option(config, bck_options)
# split the parameters
fparams = all_fpparams[:nfparams]
fobjparams = all_fpparams[nfparams:nfparams + nf_objparams]
pparams = all_fpparams[nfparams + nf_objparams:nfparams +
nf_objparams + npparams]
pobjparams = all_fpparams[nfparams + nf_objparams + npparams:]
# select the method for the sampling
if xsamples is None:
methods = {
"mh": mh,
"_dummy1d": dummy1d,
"mhcustom": mhcustom,
}
method_fcn = get_method("mcquad", methods, method)
xsamples, wsamples = method_fcn(log_pfcn, x0, pparams, **config)
epf = _integrate(ffcn, xsamples, wsamples, fparams)
# save parameters for backward calculations
ctx.xsamples = xsamples
ctx.wsamples = wsamples
ctx.ffcn = ffcn
ctx.log_pfcn = log_pfcn
ctx.fparam_sep = TensorNonTensorSeparator((*fparams, *fobjparams))
ctx.pparam_sep = TensorNonTensorSeparator((*pparams, *pobjparams))
ctx.nfparams = len(fparams)
ctx.npparams = len(pparams)
ctx.method = method
# save for backward
ftensor_params = ctx.fparam_sep.get_tensor_params()
ptensor_params = ctx.pparam_sep.get_tensor_params()
ctx.nftensorparams = len(ftensor_params)
ctx.nptensorparams = len(ptensor_params)
ctx.save_for_backward(epf, *ftensor_params, *ptensor_params)
return epf
def forward(ctx, A, B, E, M, method, fwd_options, bck_options, na,
*all_params):
# A: (*BA, nr, nr)
# B: (*BB, nr, ncols)
# E: (*BE, ncols) or None
# M: (*BM, nr, nr) or None
# all_params: list of tensor of any shape
# returns: (*BABEM, nr, ncols)
# separate the parameters for A and for M
params = all_params[:na]
mparams = all_params[na:]
config = set_default_option({}, fwd_options)
ctx.bck_config = set_default_option({}, bck_options)
if torch.all(B == 0): # special case
dims = (*_get_batchdims(A, B, E, M), *B.shape[-2:])
x = torch.zeros(dims, dtype=B.dtype, device=B.device)
else:
with A.uselinopparams(*params), M.uselinopparams(
*mparams) if M is not None else dummy_context_manager():
methods = {
"custom_exactsolve": custom_exactsolve,
"scipy_gmres": wrap_gmres,
"broyden1": broyden1_solve,
"cg": cg,
"bicgstab": bicgstab,
}
method_fcn = get_method("solve", methods, method)
x = method_fcn(A, B, E, M, **config)
ctx.e_is_none = E is None
ctx.A = A
ctx.M = M
if ctx.e_is_none:
ctx.save_for_backward(x, *all_params)
else:
ctx.save_for_backward(x, E, *all_params)
ctx.na = na
return x
def forward(ctx, A, neig, mode, M, fwd_options, bck_options, na,
*amparams):
# A: LinearOperator (*BA, q, q)
# M: LinearOperator (*BM, q, q) or None
# separate the sets of parameters
params = amparams[:na]
mparams = amparams[na:]
config = set_default_option({}, fwd_options)
ctx.bck_config = set_default_option(
{
"degen_atol": None,
"degen_rtol": None,
}, bck_options)
# options for calculating the backward (not for `solve`)
alg_keys = ["degen_atol", "degen_rtol"]
ctx.bck_alg_config = get_and_pop_keys(ctx.bck_config, alg_keys)
method = config.pop("method")
with A.uselinopparams(*params), M.uselinopparams(
*mparams) if M is not None else dummy_context_manager():
methods = {
"davidson": davidson,
"custom_exacteig": custom_exacteig,
}
method_fcn = get_method("symeig", methods, method)
evals, evecs = method_fcn(A, neig, mode, M, **config)
# save for the backward
# evals: (*BAM, neig)
# evecs: (*BAM, na, neig)
ctx.save_for_backward(evals, evecs, *amparams)
ctx.na = na
ctx.A = A
ctx.M = M
return evals, evecs
