grad_outputs=gyfcn,
create_graph=torch.is_grad_enabled())
grad_nontensor_params = [
None for _ in range(param_sep.nnontensors())
]
grad_params = param_sep.reconstruct_params(grad_tensor_params,
grad_nontensor_params)
return (None, None, None, None, None, *grad_params)
def _get_rootfinder_default_method(method):
if method is None:
return "broyden1"
else:
return method
def _get_minimizer_default_method(method):
if method is None:
return "broyden1"
else:
return method
# docstring completion
rf_methods: Sequence[Callable] = [broyden1, broyden2, linearmixing]
rootfinder.__doc__ = get_methods_docstr(rootfinder, rf_methods)
equilibrium.__doc__ = get_methods_docstr(equilibrium, rf_methods)
minimize.__doc__ = get_methods_docstr(minimize, rf_methods)
**ctx.bck_config)
# only take the output for the earliest time
states = [out[-1] for out in outs]
states[y_index] = yt[t_flip_idx]
# gyt is the contribution from the input grad_y
# gy0 is the propagated gradients from the later time step
states[dLdy_index] = grad_yt[t_flip_idx] + states[dLdy_index]
if ts_requires_grad:
grad_ts[0] = states[dLdt_index].reshape(-1)
grad_y0 = states[dLdy_index] # dL/dy0, (*ny)
if ts_requires_grad:
grad_ts = torch.cat(grad_ts).reshape(*ts.shape)
grad_tensor_params = states[dLdp_slice]
grad_ntensor_params = [
None for _ in range(len(allparams) - ntensor_params)
]
grad_params = param_sep.reconstruct_params(grad_tensor_params,
grad_ntensor_params)
return (None, grad_ts, None, None, None, grad_y0, *grad_params)
# docstring completion
ivp_methods: Dict[str, Callable] = {
"rk45": rk45_adaptive,
"rk23": rk23_adaptive,
"rk4": rk4_ivp,
}
solve_ivp.__doc__ = get_methods_docstr(solve_ivp, ivp_methods)
Returns
-------
torch.Tensor
The integrated values.
"""
swapaxes = dim != -1
if swapaxes:
y = y.transpose(dim, -1)
if y.shape[-1] != self.nx:
raise RuntimeError("The length of integrated dimension does not match with x")
res = self.obj.integrate(y)
if keepdim:
res = res.unsqueeze(-1)
if swapaxes:
res = res.transpose(dim, -1)
return res
def getparamnames(self, methodname: str, prefix: str = "") -> List[str]:
""""""
return self.obj.getparamnames(methodname, prefix=prefix + "obj.")
# docstring completion
_squad_methods = {
"cspline": CubicSplineSQuad,
# "simpson": SimpsonSQuad,
"trapz": TrapzSQuad,
}
SQuad.__doc__ = get_methods_docstr(SQuad, _squad_methods)
The values at the given position with shape ``(*BY, nr)``.
If ``y`` has been specified during ``__init__`` and also
specified here, the value of ``y`` given here will be ignored.
If no ``y`` ever specified, then it will raise an error.
Returns
-------
torch.Tensor
The interpolated values with shape ``(*BY, nrq)``.
"""
def __init__(self, x, y=None, method=None, **fwd_options):
if method is None:
method = "cspline"
if method == "cspline":
self.obj = CubicSpline1D(x, y, **fwd_options)
else:
raise RuntimeError("Unknown interp1d method: %s" % method)
def __call__(self, xq, y=None):
return self.obj(xq, y)
def getparamnames(self, methodname, prefix=""):
return [prefix + "obj." + c for c in self.obj.getparamnames()]
# docstring completion
interp1d_methods = {
"cspline": CubicSpline1D,
}
Interp1D.__doc__ = get_methods_docstr(Interp1D, interp1d_methods)
grad_mparams = torch.autograd.grad(
(mloss, ),
mparams,
grad_outputs=(v, ),
create_graph=torch.is_grad_enabled(),
allow_unused=True)
return (None, grad_B, grad_E, None, None, None, None, None,
*grad_params, *grad_mparams)
def custom_exactsolve(A, B, E=None, M=None, **options):
# A: (*BA, na, na)
# B: (*BB, na, ncols)
# E: (*BE, ncols)
# M: (*BM, na, na)
return exactsolve(A, B, E, M)
# docstring completion
_solve_methods = {
"cg": cg,
"bicgstab": bicgstab,
"exactsolve": exactsolve,
"broyden1": broyden1_solve,
"scipy_gmres": wrap_gmres,
}
ignore_kwargs = ["E", "M", "mparams"]
solve.__doc__ = get_methods_docstr(solve, _solve_methods, ignore_kwargs)
tensor_params_copy,
grad_outputs=gyfcn,
create_graph=torch.is_grad_enabled())
grad_nontensor_params = [
None for _ in range(param_sep.nnontensors())
]
grad_params = param_sep.reconstruct_params(grad_tensor_params,
grad_nontensor_params)
return (None, None, None, None, None, *grad_params)
def _get_rootfinder_default_method(method):
if method is None:
return "broyden1"
else:
return method
def _get_minimizer_default_method(method):
if method is None:
return "broyden1"
else:
return method
# docstring completion
rootfinder.__doc__ = get_methods_docstr(rootfinder, [broyden1])
equilibrium.__doc__ = get_methods_docstr(equilibrium, [broyden1])
minimize.__doc__ = get_methods_docstr(minimize, [broyden1])
log_pfcn,
x0=xsamples[0], # unused because xsamples is set
xsamples=xsamples,
wsamples=wsamples,
fparams=(grad_epf, epf, *fptensor_params_copy),
pparams=pparams,
fwd_options=ctx.bck_config,
bck_options=ctx.bck_config)
dLdthetaf = aug_epfs[:nftensorparams]
dLdthetap = aug_epfs[nftensorparams:]
# combine the gradient for all fparams
dLdfnontensor = [None for _ in range(ctx.fparam_sep.nnontensors())]
dLdpnontensor = [None for _ in range(ctx.pparam_sep.nnontensors())]
dLdtf = ctx.fparam_sep.reconstruct_params(dLdthetaf, dLdfnontensor)
dLdtp = ctx.pparam_sep.reconstruct_params(dLdthetap, dLdpnontensor)
return (None, None, None, None, None, None, None, None, None, None,
*dLdtf, *dLdtp)
def _integrate(ffcn, xsamples, wsamples, fparams):
nsamples = len(xsamples)
res = 0.0
for x, w in zip(xsamples, wsamples):
res = res + ffcn(x, *fparams) * w
return res
# docstring completion
mcquad.__doc__ = get_methods_docstr(mcquad, [mh, mhcustom])
# the contribution from the parallel elements
gevecsM_par = (-0.5 *
torch.einsum("...ae,...ae->...e", grad_evecs, evecs)
).unsqueeze(-2) * evecs # (*BAM, na, neig)
gaccumM = gevalsM + gevecsM + gevecsM_par
grad_mparams = torch.autograd.grad(
outputs=(mloss, ),
inputs=mparams,
grad_outputs=(gaccumM, ),
create_graph=torch.is_grad_enabled(),
)
return (None, None, None, None, None, None, None, *grad_params,
*grad_mparams)
def custom_exacteig(A, neig, mode, M=None, **options):
return exacteig(A, neig, mode, M)
# docstring completion
_symeig_methods = {
"exacteig": exacteig,
"davidson": davidson,
}
ignore_kwargs = ["M", "mparams"]
symeig.__doc__ = get_methods_docstr(symeig, _symeig_methods, ignore_kwargs)
svd.__doc__ = get_methods_docstr(svd, _symeig_methods)
class _BaseInfTransform(object):
@abstractmethod
def forward(self, t):
pass
@abstractmethod
def dxdt(self, t):
pass
@abstractmethod
def x2t(self, x):
pass
class _TanInfTransform(_BaseInfTransform):
def forward(self, t):
return torch.tan(t)
def dxdt(self, t):
sec = 1. / torch.cos(t)
return sec * sec
def x2t(self, x):
return torch.atan(x)
# docstring completion
quad.__doc__ = get_methods_docstr(quad, [leggauss])