xitorch/xitorch

Differentiable scientific computing library

xitorch: differentiable scientific computing library

xitorch is a PyTorch-based library of differentiable functions and functionals that can be widely used in scientific computing applications as well as deep learning.

The documentation can be found at: https://xitorch.readthedocs.io/

Example

Finding root of a function:

import torch
from xitorch.optimize import rootfinder

def func1(y, A):  # example function
    return torch.tanh(A @ y + 0.1) + y / 2.0

# set up the parameters and the initial guess
A = torch.tensor([[1.1, 0.4], [0.3, 0.8]]).requires_grad_()
y0 = torch.zeros((2, 1))  # zeros as the initial guess

# finding a root
yroot = rootfinder(func1, y0, params=(A,))

# calculate the derivatives
dydA, = torch.autograd.grad(yroot.sum(), (A,), create_graph=True)
grad2A, = torch.autograd.grad(dydA.sum(), (A,), create_graph=True)

Modules

• linalg: Linear algebra and sparse linear algebra module
• optimize: Optimization and root finder module
• integrate: Quadrature and integration module
• interpolate: Interpolation

Requirements

• python >=3.8.1,<3.12
• pytorch 1.13.1 or higher (install here)

Getting started

After fulfilling all the requirements, type the commands below to install xitorch

python -m pip install xitorch

Or to install from GitHub:

python -m pip install git+https://github.com/xitorch/xitorch.git

Finally, if you want to make an editable install from source:

git clone https://github.com/xitorch/xitorch.git
cd xitorch
python -m pip install -e .

Note that the last option is only available per PEP 660, so you will require pip >= 23.1

Used in

• Differentiable Quantum Chemistry (DQC): https://dqc.readthedocs.io/

Gallery

Neural mirror design (example 01):

Initial velocity optimization in molecular dynamics (example 02):

xitorch's Issues

Help using rootfinder

Hello,
rootfinder solves equation 0 = f(y,\theta). This equation may have no solution. Does it return a value, or a parameter that I can use to detect this is no solution ?
Thanks !

Cannot install from sdist obtained from PyPI

Describe the bug

Installation with pip from sdist fails (version 0.3.0).

To Reproduce

Presumably, the directory layout from the git repository is assumed in setup.py, but the relevant files are not actually included in the sdist.

source tree in: /home/conda/staged-recipes/build_artifacts/xitorch_1649528643673/work
export PREFIX=/home/conda/staged-recipes/build_artifacts/xitorch_1649528643673/_h_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_pla
export BUILD_PREFIX=/home/conda/staged-recipes/build_artifacts/xitorch_1649528643673/_build_env
export SRC_DIR=/home/conda/staged-recipes/build_artifacts/xitorch_1649528643673/work
Using pip 22.0.4 from $PREFIX/lib/python3.10/site-packages/pip (python 3.10)
Non-user install because user site-packages disabled
Ignoring indexes: https://pypi.org/simple
Created temporary directory: /tmp/pip-ephem-wheel-cache-aj1o1cc3
Created temporary directory: /tmp/pip-req-tracker-hqkvtmg0
Initialized build tracking at /tmp/pip-req-tracker-hqkvtmg0
Created build tracker: /tmp/pip-req-tracker-hqkvtmg0
Entered build tracker: /tmp/pip-req-tracker-hqkvtmg0
Created temporary directory: /tmp/pip-install-gylksics
Processing $SRC_DIR
  Added file://$SRC_DIR to build tracker '/tmp/pip-req-tracker-hqkvtmg0'
  Running setup.py (path:$SRC_DIR/setup.py) egg_info for package from file://$SRC_DIR
  Created temporary directory: /tmp/pip-pip-egg-info-uaj5m5_v
  Running command python setup.py egg_info
  Preparing metadata (setup.py): started
  Traceback (most recent call last):
    File "<string>", line 2, in <module>
    File "<pip-setuptools-caller>", line 34, in <module>
    File "/home/conda/staged-recipes/build_artifacts/xitorch_1649528643673/work/setup.py", line 49, in <module>
      install_requires=get_requirements("requirements.txt"),
    File "/home/conda/staged-recipes/build_artifacts/xitorch_1649528643673/work/setup.py", line 32, in get_requirements
      with open(absdir(fname), "r") as f:
  FileNotFoundError: [Errno 2] No such file or directory: '/home/conda/staged-recipes/build_artifacts/xitorch_1649528643673/work/requirements.txt'
  error: subprocess-exited-with-error
  
  × python setup.py egg_info did not run successfully.
  │ exit code: 1
  ╰─> See above for output.
  
  note: This error originates from a subprocess, and is likely not a problem with pip.
  full command: /home/conda/staged-recipes/build_artifacts/xitorch_1649528643673/_h_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_pla/bin/python -c '
  exec(compile('"'"''"'"''"'"'
  # This is <pip-setuptools-caller> -- a caller that pip uses to run setup.py
  #
  # - It imports setuptools before invoking setup.py, to enable projects that directly
  #   import from `distutils.core` to work with newer packaging standards.
  # - It provides a clear error message when setuptools is not installed.
  # - It sets `sys.argv[0]` to the underlying `setup.py`, when invoking `setup.py` so
  #   setuptools doesn'"'"'t think the script is `-c`. This avoids the following warning:
  #     manifest_maker: standard file '"'"'-c'"'"' not found".
  # - It generates a shim setup.py, for handling setup.cfg-only projects.
  import os, sys, tokenize

  try:
      import setuptools
  except ImportError as error:
      print(
          "ERROR: Can not execute `setup.py` since setuptools is not available in "
          "the build environment.",
          file=sys.stderr,
      )
      sys.exit(1)
  
  __file__ = %r
  sys.argv[0] = __file__
  
  if os.path.exists(__file__):
      filename = __file__
      with tokenize.open(__file__) as f:
          setup_py_code = f.read()
  else:
      filename = "<auto-generated setuptools caller>"
      setup_py_code = "from setuptools import setup; setup()"
  
  exec(compile(setup_py_code, filename, "exec"))
  '"'"''"'"''"'"' % ('"'"'/home/conda/staged-recipes/build_artifacts/xitorch_1649528643673/work/setup.py'"'"',), "<pip-setuptools-caller>", "exec"))' egg_info --egg-base /tmp/pip-pip-egg-info-uaj5m5_v
  cwd: /home/conda/staged-recipes/build_artifacts/xitorch_1649528643673/work/
  Preparing metadata (setup.py): finished with status 'error'
error: metadata-generation-failed

× Encountered error while generating package metadata.
╰─> See above for output.

Expected behavior

Install cleanly from sdist.

Systems:

• OS: Linux, MacOS, Windows
• Python version: 3.10
• PyTorch version: 1.11
• xitorch version: 0.3.0

Additional context

See conda-forge/staged-recipes#18636

`torch.outer` is not yet present in pytorch

Dear developer,

You are using a function torch.outer that is not available in pytorch 1.6.0.

Custom termination criterion for nonlinear solver

Is your feature request related to a problem? Please describe.
When using Xitorch through DQC, termination criteria for the Kohn-Sham solver are given in terms of tolerances on the self-consistent parameter. However, in most research circumstances, tolerances on the energy of the system are a more appropriate measure. These cannot currently be implemented in DQC, and the change to accommodate them has to be made in Xitorch.

Describe the solution you'd like
A keyword argument custom_terminator passed to xitorch._impls.optimize.root.rootsolver._nonlin_solver through the fwd_options dictionary from DQC's KS.run(...).

The value passed into the keyword argument should be an object which defines a method check(x: torch.Tensor, y: torch.Tensor, dx: torch.Tensor) -> bool, and evaluates to True if the termination criterion has been met (otherwise False).

If no value is passed, this should default to xitorch._impls.optimize.root.rootsolver.TerminationCondition((f_tol, f_rtol, y_norm, x_tol, x_rtol) which is the current criterion. The signature of the check method on TerminationCondition is modified from check(xnorm: float, ynorm: float, dxnorm: float) -> bool to the more generic signature given above, for consistency.

Describe alternatives you've considered

• An external criterion varying x_tol, and f_tol depending on the energy of the current solution → inaccurate and requires significant changes to DQC
• Adding an energy termination criterion to Xitorch and providing a boolean flag to determine which criterion should be used → less flexible, cannot provide a history of the convergence energies

Overall, this seems to be a smooth solution, providing maximal flexibility for minimal changes, and expanding on the scientific capabilities of both Xitorch and DQC. No other alternatives considered come close to these effects.

Additional context
This change ensures we can use DQC and Xitorch with significant speed-up without compromising accuracy. I can implement it myself, and would ideally have it approved before Christmas.

Optimize module to-do list

• rootfinder and equilibrium (can take a look at SciPy, good first issue):
  • (method="broyden1")
  • (method="broyden2")
  • (method="anderson")
  • (method="diagbroyden")
  • (method="linearmixing")
  • (method="excitingmixing")
  • (method="krylov")
• minimize
  • (method="gd") (vanilla gradient descent)
  • (method="lbfgs")
  • (method="cg")
• mincon (constrained optimization: min f(x,p) s.t. g(x,q) = 0)

gradient-vector product and vector-gradient product for higher-order gradients

Hi,

I'm trying to see what's the best way to calculate gradient-vector product (gvp) and vector-gradient product (vgp) for hessian, and third-order gradients?

Thanks,
JW

Example 2 broken

Running the second example 02, I run into following error:

Traceback (most recent call last):
File "/home/jwittke/py_tmp/main.py", line 125, in
mainopt()
File "/home/jwittke/py_tmp/main.py", line 110, in mainopt
loss, yt = get_loss(pos, vel, ts, pos_target)
File "/home/jwittke/py_tmp/main.py", line 29, in get_loss
yt = solve_ivp(dydt, ts, y0, fwd_options={"method": "rk4"})
File "/home/jwittke/anaconda3/lib/python3.7/site-packages/xitorch/integrate/solve_ivp.py", line 92, in solve_ivp
return _SolveIVP.apply(pfcn, ts, fwd_options, bck_options, len(params), y0, *params, *pfcn.objparams())
File "/home/jwittke/anaconda3/lib/python3.7/site-packages/xitorch/integrate/solve_ivp.py", line 114, in forward
yt = solver(pfcn, ts, y0, params, **config)
File "/home/jwittke/anaconda3/lib/python3.7/site-packages/xitorch/_impls/integrate/ivp/adaptive_rk.py", line 178, in rk45_adaptive
return _rk_adaptive(fcn, ts, y0, params, RK45, **kwargs)
File "/home/jwittke/anaconda3/lib/python3.7/site-packages/xitorch/_impls/integrate/ivp/adaptive_rk.py", line 164, in _rk_adaptive
return solver.solve()
File "/home/jwittke/anaconda3/lib/python3.7/site-packages/xitorch/_impls/integrate/ivp/adaptive_rk.py", line 72, in solve
rk_state = self._step(rk_state, ts[i])
File "/home/jwittke/anaconda3/lib/python3.7/site-packages/xitorch/_impls/integrate/ivp/adaptive_rk.py", line 83, in _step
rk_state, t1_achieved = self._single_step(rk_state, t1)
File "/home/jwittke/anaconda3/lib/python3.7/site-packages/xitorch/_impls/integrate/ivp/adaptive_rk.py", line 98, in _single_step
ynew, fnew = rk_step(self.func, t0, y0, f0, hstep, abck)
File "/home/jwittke/anaconda3/lib/python3.7/site-packages/xitorch/_impls/integrate/ivp/adaptive_rk.py", line 16, in rk_step
K[s] = func(t + c * h, y + dy)
RuntimeError: The size of tensor a (2) must match the size of tensor b (128) at non-singleton dimension 3

I am using:

conda version : 4.9.2
conda-build version : 3.18.11
python version : 3.7.6.final.0
platform : linux-64
user-agent : conda/4.9.2 requests/2.23.0 CPython/3.7.6 Linux/4.4.0-194-generic ubuntu/16.04.7 glibc/2.23

pytorch 1.7.1 py3.7_cpu_0 [cpuonly] pytorch

A RuntimeError occur when call _Jac.fullmatrix()

This occurred when I was implementing an iterative algorithm and calculating gradients. With a call:
grad = jac(...), grad.fullmatrix(). But it does not occur when I use grad.H.fullmatrix()

I tried to print other values of grad, like grad.dfdy, grad.yout, and it works just fine. As the error information says, it might be the gradient of dfdy with respect to grad.v is lost.

File "/home/tju/zyzh/abinitialTransport/calc/SCF.py", line 182, in backward
grad.fullmatrix()
File "/home/tju/.conda/envs/abinit/lib/python3.8/site-packages/xitorch/_core/linop.py", line 354, in fullmatrix
return self.mm(V) # (B1,B2,...,Bb,np,nq)
File "/home/tju/.conda/envs/abinit/lib/python3.8/site-packages/xitorch/_core/linop.py", line 272, in mm
ynew = self._mv(xnew) # (r,...,p)
File "/home/tju/.conda/envs/abinit/lib/python3.8/site-packages/xitorch/grad/jachess.py", line 167, in _mv
dfdyf, = torch.autograd.grad(dfdy, (v,), grad_outputs=gy1[i].reshape(self.inshape),
File "/home/tju/.conda/envs/abinit/lib/python3.8/site-packages/torch/autograd/init.py", line 226, in grad
return Variable._execution_engine.run_backward(
RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.

To Reproduce
This is the function I wrote.

class SCF(torch.autograd.Function):
@staticmethod
def forward(ctx, fcn, x0, maxIter, err, method='default', *params):
# with torch.no_grad():
# x = fcn(x0, *params)
x_ = fcn(x0, *params)

    if method == "default":
        it = 0
        old_x = x0
        while (x_-old_x).norm() > err and it < maxIter:
            old_x = x_
            x_ = fcn(x_, *params)

    elif method == 'LBFGS':
        optim = LBFGS(params=[x_], tolerance_grad=1e-10, tolerance_change=1e-15)

        def new_fcn():
            return (x_ - fcn(x_, *params)).abs().sum()

        for i in range(maxIter):
            optim.step(new_fcn)

    ctx.save_for_backward(x_, *params)
    ctx.fcn = fcn

    return x_



@staticmethod
def backward(ctx, grad_outputs):
    x_ = ctx.saved_tensors[0].clone().requires_grad_()
    params = ctx.saved_tensors[1:]

    idx = [i for i in range(len(params)) if params[i].requires_grad]


    fcn = ctx.fcn
    def new_fcn(x, *params):
        return - fcn(x, *params)

    grad = jac(fcn=new_fcn, params=(x_, *params), idxs=[0])[0]
    grad.fullmatrix()

    pre = tLA.solve(grad.H.fullmatrix().real, -grad_outputs.reshape(-1, 1))
    pre = pre.reshape(grad_outputs.shape)


    with torch.enable_grad():
        params_copy = [p.clone().requires_grad_() for p in params]
        yfcn = new_fcn(x_, *params_copy)

    grad = torch.autograd.grad(yfcn, [params_copy[i] for i in idx], grad_outputs=pre,
                               create_graph=torch.is_grad_enabled(),
                               allow_unused=True)
    grad_out = [None for _ in range(len(params))]
    for i in range(len(idx)):
        grad_out[idx[i]] = grad[i]


    return None, None, None, None, None, *grad_out

Expected behavior
Much thanks if there are any helps to find out what is the problem.

Linear algebra module to-do list

• svd (taking the n lowest or uppermost decomposition, just like symeig)
• solve
  • (method="cg")
  • (method="bicgstab")
  • (method="gmres") (#10)
  • (method="lgmres")
• symeig (eigen decomposition for symmetric matrix)
  • (method="lobpcg")
  • Stabilizing the backward calculation for degenerate case
• lstsq (minimizing |AX - B|^2)
• Minimizing |X|^2 s.t. |AX-B| for overdetermined problem

Using root finder to solve a inverse function with the NN form Error

Thank you for the great job. I find it very helpful for me to solve the inverse function. However, I have a problem as follows:

1. I want to solve a inverse function with the form of a neural network, numerically. So I choose to use root finding algorithm.
2. Following the instructions in your document, here is my codes:

torch_model = torch.nn.Sequential(
    torch.nn.Linear(3, 1),
    torch.nn.ReLU(),
)

def func1(y, A):  # example function 
    return A - torch_model(y)

A =  10*torch.ones((2,1)).requires_grad_() # output shape of the torch_model_1 is (2,1)
y0 = torch.zeros((2,3)) # the torch_model accepts input shape of (2,3)

# finding a root
yroot = rootfinder(func1, y0, params=(A,))

The torch_model accepts input shape of (batch_size,3) and the outputs (batch_size, 1)
Given the outputs, I want to find the inputs values. So, I choose to let the param A to be the given outputs and try to get the yroot.
But the code show the errors:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[42], line 9
      6 y0 = torch.zeros((2,3)) # 
      8 # finding a root
----> 9 yroot = rootfinder(func1, y0, params=(A,))

File [~/anaconda3/envs/dpc/lib/python3.8/site-packages/xitorch/optimize/rootfinder.py:93](https://file+.vscode-resource.vscode-cdn.net/media/lk/disk1/DL_lk/202205_RZ_GAN%E6%A8%A1%E5%9E%8B/model/EvoOpt/~/anaconda3/envs/dpc/lib/python3.8/site-packages/xitorch/optimize/rootfinder.py:93), in rootfinder(fcn, y0, params, bck_options, method, **fwd_options)
     91 pfunc = get_pure_function(fcn)
     92 fwd_options["method"] = _get_rootfinder_default_method(method)
---> 93 return _RootFinder.apply(pfunc, y0, pfunc, False, fwd_options, bck_options,
     94                          len(params), *params, *pfunc.objparams())

File [~/anaconda3/envs/dpc/lib/python3.8/site-packages/torch/autograd/function.py:506](https://file+.vscode-resource.vscode-cdn.net/media/lk/disk1/DL_lk/202205_RZ_GAN%E6%A8%A1%E5%9E%8B/model/EvoOpt/~/anaconda3/envs/dpc/lib/python3.8/site-packages/torch/autograd/function.py:506), in Function.apply(cls, *args, **kwargs)
    503 if not torch._C._are_functorch_transforms_active():
    504     # See NOTE: [functorch vjp and autograd interaction]
    505     args = _functorch.utils.unwrap_dead_wrappers(args)
--> 506     return super().apply(*args, **kwargs)  # type: ignore[misc]
    508 if cls.setup_context == _SingleLevelFunction.setup_context:
    509     raise RuntimeError(
    510         'In order to use an autograd.Function with functorch transforms '
    511         '(vmap, grad, jvp, jacrev, ...), it must override the setup_context '
    512         'staticmethod. For more details, please see '
    513         'https://pytorch.org/docs/master/notes/extending.func.html')

File [~/anaconda3/envs/dpc/lib/python3.8/site-packages/xitorch/optimize/rootfinder.py:302](https://file+.vscode-resource.vscode-cdn.net/media/lk/disk1/DL_lk/202205_RZ_GAN%E6%A8%A1%E5%9E%8B/model/EvoOpt/~/anaconda3/envs/dpc/lib/python3.8/site-packages/xitorch/optimize/rootfinder.py:302), in _RootFinder.forward(ctx, fcn, y0, fwd_fcn, is_opt_method, options, bck_options, nparams, *allparams)
    300     name = "rootfinder" if not is_opt_method else "minimizer"
    301     method_fcn = get_method(name, methods, method)
--> 302     y = method_fcn(fwd_fcn, y0, params, **config)
    304 ctx.fcn = fcn
    305 ctx.is_opt_method = is_opt_method

File [~/anaconda3/envs/dpc/lib/python3.8/site-packages/xitorch/_impls/optimize/root/rootsolver.py:181](https://file+.vscode-resource.vscode-cdn.net/media/lk/disk1/DL_lk/202205_RZ_GAN%E6%A8%A1%E5%9E%8B/model/EvoOpt/~/anaconda3/envs/dpc/lib/python3.8/site-packages/xitorch/_impls/optimize/root/rootsolver.py:181), in broyden1(fcn, x0, params, **kwargs)
    167 @functools.wraps(_nonlin_solver, assigned=('__annotations__',))  # takes only the signature
    168 def broyden1(fcn, x0, params=(), **kwargs):
    169     """
    170     Solve the root finder or linear equation using the first Broyden method [1]_.
    171     It can be used to solve minimization by finding the root of the
   (...)
    179            https://web.archive.org/web/20161022015821/http://www.math.leidenuniv.nl/scripties/Rotten.pdf
    180     """
--> 181     return _nonlin_solver(fcn, x0, params, "broyden1", **kwargs)

File [~/anaconda3/envs/dpc/lib/python3.8/site-packages/xitorch/_impls/optimize/root/rootsolver.py:123](https://file+.vscode-resource.vscode-cdn.net/media/lk/disk1/DL_lk/202205_RZ_GAN%E6%A8%A1%E5%9E%8B/model/EvoOpt/~/anaconda3/envs/dpc/lib/python3.8/site-packages/xitorch/_impls/optimize/root/rootsolver.py:123), in _nonlin_solver(fcn, x0, params, method, alpha, uv0, max_rank, maxiter, f_tol, f_rtol, x_tol, x_rtol, line_search, verbose, **unused)
    118     raise ValueError("Jacobian inversion yielded zero vector. "
    119                      "This indicates a bug in the Jacobian "
    120                      "approximation.")
    122 if line_search:
--> 123     s, xnew, ynew, y_norm_new = _nonline_line_search(func, x, y, dx,
    124                                                      search_type=line_search)
    125 else:
    126     s = 1.0

File [~/anaconda3/envs/dpc/lib/python3.8/site-packages/xitorch/_impls/optimize/root/rootsolver.py:277](https://file+.vscode-resource.vscode-cdn.net/media/lk/disk1/DL_lk/202205_RZ_GAN%E6%A8%A1%E5%9E%8B/model/EvoOpt/~/anaconda3/envs/dpc/lib/python3.8/site-packages/xitorch/_impls/optimize/root/rootsolver.py:277), in _nonline_line_search(func, x, y, dx, search_type, rdiff, smin)
    274     return (phi(s + ds, store=False) - phi(s)) [/](https://file+.vscode-resource.vscode-cdn.net/) ds
    276 if search_type == 'armijo':
--> 277     s, phi1 = _scalar_search_armijo(phi, tmp_phi[0], -tmp_phi[0],
    278                                     amin=smin)
    280 if s is None:
    281     # No suitable step length found. Take the full Newton step,
    282     # and hope for the best.
    283     s = 1.0

File [~/anaconda3/envs/dpc/lib/python3.8/site-packages/xitorch/_impls/optimize/root/rootsolver.py:295](https://file+.vscode-resource.vscode-cdn.net/media/lk/disk1/DL_lk/202205_RZ_GAN%E6%A8%A1%E5%9E%8B/model/EvoOpt/~/anaconda3/envs/dpc/lib/python3.8/site-packages/xitorch/_impls/optimize/root/rootsolver.py:295), in _scalar_search_armijo(phi, phi0, derphi0, c1, alpha0, amin, max_niter)
    294 def _scalar_search_armijo(phi, phi0, derphi0, c1=1e-4, alpha0=1, amin=0, max_niter=20):
--> 295     phi_a0 = phi(alpha0)
    296     if phi_a0 <= phi0 + c1 * alpha0 * derphi0:
    297         return alpha0, phi_a0

File [~/anaconda3/envs/dpc/lib/python3.8/site-packages/xitorch/_impls/optimize/root/rootsolver.py:263](https://file+.vscode-resource.vscode-cdn.net/media/lk/disk1/DL_lk/202205_RZ_GAN%E6%A8%A1%E5%9E%8B/model/EvoOpt/~/anaconda3/envs/dpc/lib/python3.8/site-packages/xitorch/_impls/optimize/root/rootsolver.py:263), in _nonline_line_search..phi(s, store)
    261 if s == tmp_s[0]:
    262     return tmp_phi[0]
--> 263 xt = x + s * dx
    264 v = func(xt)
    265 p = _safe_norm(v)**2

RuntimeError: The size of tensor a (6) must match the size of tensor b (2) at non-singleton dimension 0

Could anyone to give some suggestions? Thanks!

Rootfinder fails with a simple case

Describe the bug
xitorch.optimize.rootfinder fails with a simple case.

To Reproduce

import torch
from xitorch.optimize import rootfinder

def fcn(x: torch.Tensor, r: torch.Tensor) -> torch.Tensor:
    res0 = x[0] - x[1] - 1.0
    res1 = x[2] + x[3]
    res2 = x[0] - x[1] - x[3] * r
    res3 = x[4] + x[5]
    res4 = x[2] + x[4]
    res5 = x[1]
    return torch.stack([res2, res3, res0, res1, res4, res5])

x0 = torch.zeros(6)
r = torch.tensor(1.0)
x = rootfinder(fcn, x0, params=(r,))
print(x)

Produces

/mnt/c/Users/firma/Documents/Projects/Git/xitorch/xitorch/_impls/optimize/root/rootsolver.py:163: ConvergenceWarning: The rootfinder does not converge after 700 iterations. Best |dx|=0.000e+00, |f|=1.000e+00 at iter 0
  warnings.warn(ConvergenceWarning(msg))

Expected behavior
This should not fail. Using numpy and scipy.optimize.root succeed on this case (even with the same broyden1 method).

Systems:

• OS: WSL2 Ubuntu
• Python version: 3.9.7
• PyTorch version: 1.10.0
• xitorch version: 0.4.0.dev0+1e875cc

Additional context
Changing the res order from torch.stack([res2, res3, res0, res1, res4, res5]) to torch.stack([res0, res1, res2, res3, res4, res5]) makes it successful.

Interpolate module to-do list

• gridinterp1 (interpolation where the input points are given on a grid)
  • method="cspline" (cubic spline)
  • method="linear"
• interp1 (interpolation where input points are not necessarily on a grid)
  • method="cspline"
  • method="linear"
• gridinterp2
  • method="bilinear"
  • method="rectbivariatespline"
• gridinterp3
  • method="trilinear"
• Unstructured N-dimensional data interpolation

Status Badges issues

Describe the bug
The status badges are bugged and require fixing / updating

To Reproduce
Steps to reproduce the behavior:

1. Look at the status badges in README.md.

Expected behavior
The status badges mirror the tests, which are successful

Systems:

• OS: Any
• Python version: 3.8.12
• PyTorch version: 1.13.1
• xitorch version: 0.4.1

Incompatibility with pytorch 1.10

Describe the bug
Upgrading the pytorch to 1.10 makes some of the test fails (while it succeeds in 1.9 & 1.8).

To Reproduce
Steps to reproduce the behavior:

1. Install xitorch with pytorch 1.10
2. Run the test (cd xitorch/_tests/; pytest)
3. See error

FAILED test_optimize.py::test_equil[dtype0-device0-DummyModule] - torch.autograd.gradcheck.GradcheckError: Jacobian m...
FAILED test_optimize.py::test_equil[dtype1-device1-DummyNNModule] - torch.autograd.gradcheck.GradcheckError: Jacobian...
FAILED test_optimize.py::test_equil[dtype2-device2-DummyModule] - torch.autograd.gradcheck.GradcheckError: While cons...

Expected behavior
Seamless transition to 1.10

Systems:

• OS: Ubuntu 20.04 WSL
• Python version: 3.8.5
• PyTorch version: 1.10
• xitorch version: 0.4.0.dev0+47631ef

Additional context

Joint eigenvalues of a pair of matrix

Is your feature request related to a problem? Please describe.
In a nutshell, I would like to have a differentiable function for joint eigenvalues computation, an MATLAB interface can be found here https://stackoverflow.com/questions/36551182/how-can-i-find-the-joint-eigenvalues-of-two-matrices-in-matlab.

Describe the solution you'd like
Specifically, I would like to use a distance of covariance matrices as a loss function.

Upon my research, I think the torch.lobpcg() function can do that, but only for k largest, and the back propagation does not work due to this thread which you are also involved in (pytorch/pytorch#38948).

Describe alternatives you've considered
I'm way too layman to solve algebra like this. But based on my research, I've located it as generalized eigenvalue problem (http://fourier.eng.hmc.edu/e161/lectures/algebra/node7.html). I also find this work around (https://www.alglib.net/eigen/symmetric/generalizedsymmevd.php), do you have any idea on how to make such joint eigenvalue solver differentiable?

Additional context
Add any other context or screenshots about the feature request here.

Complex hermitian matrices not recognized as hermitian.

Describe the bug
We were using xitorch with complex hermitian matrices and noticed that it would not check for hermiticity but for symmetry, therefore it raised a wrong error message.
is_hermitian = torch.allclose(mat, mat.transpose(-2, -1))
Of course this can be fixed by adding a .conj() (tested it error message disappears), however I am not sure if anything else in the code relies on the matrix being real.

To Reproduce

import torch, xitorch
from xitorch import linalg 

a = torch.rand(2,2) + 1j*torch.rand(2,2)
a = a+a.T.conj()
op = xitorch.LinearOperator.m(a)
linalg.symeig(op)

Expected behavior
Complex hermitian operators are recognized as such.
Systems:

• Python version: 3.88
• PyTorch version: 1.9
• xitorch version: 0.2.0

Integrate module to-do list

• solve_ivp
  • (method="rk45") (RK4(5) with adaptive step size)
  • (method="dop853") (can take a look at SciPy, good first issue)
  • (method="bdf")
  • (method="lsoda")
  • (method="cvode")
• quad
  • error-based integration (like SciPy's quad)
• mcquad
  • (method="nuts")
  • (method="mh")
  • (method="custom_mh") (Metropolis-Hastings with custom step algorithm, e.g. for 2D Ising model)
• solve_bvp

Any way to improve the computational speed of the linalg.symeig() function?

Is your feature request related to a problem? Please describe.
Yes. My general goal is to inversely determine the material parameters of a finite element model. The algorithm includes solving the eigenproblem posed by the mass and stiffness matrices of the finite element model. The overall steps are something like

1. material_parameters = neural_network_model(inputs)
2. Kmatrix, Mmatrix = assemble(mesh, material_parameters)
3. eigen_solution = xitorch.linalg.symeig(A=Kmat, neig=100, M=Mmat)
4. loss = loss_function(eigen_solution)
5. loss.backward()

The above algorithm with xitorch works flawlessly. The only issue is that once the problem size is non-trivial (with more than 10,000 degrees of freedom), it takes too long to solve the eigenproblem (~ 3min) that training becomes impractical. I noticed that the solution speed is independent of the number of requested eigenpairs.

Describe the solution you'd like
I wonder if it's possible to speed up the eigen solution since I only need a few of the lowest eigenpairs. The matrices in the finite element system are very sparse. However, I noticed that they are converted back when used to solve the eigenproblem.

Describe alternatives you've considered
Right now, the only eigen solution package that I found that supports both general eigenproblem and backward is the xitorch.linalg.symeig() (so many thanks for this amazing work!). The eigen solver in the scipy package is fast for a subset of eigenpairs, but it doesn't support backward().

Additional context
I'm not familiar with the eigen solution algorithms, but any suggestion is appreciated. And thank you again for this amazing work!

Spurious results if x input is not sorted in interpolate.Interp1D

xitorch.interpolate.Interp1D appears to give spurious results if the x input is not sorted. This is however stated nowhere in the documentation.