summaryrefslogtreecommitdiff
path: root/numpy/dual.py
blob: c47f8f8204b35076dc4148d6d10258258555bf87 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
# This module should be used for functions both in numpy and scipy if
#  you want to use the numpy version if available but the scipy version
#  otherwise.
#  Usage  --- from numpy.dual import fft, inv

__all__ = ['fft','ifft','fftn','ifftn','fft2','ifft2',
           'norm','inv','svd','solve','det','eig','eigvals',
           'eigh','eigvalsh','lstsq', 'pinv','cholesky','i0']

import numpy.linalg as linpkg
import numpy.fft as fftpkg
from numpy.lib import i0
import sys


fft = fftpkg.fft
ifft = fftpkg.ifft
fftn = fftpkg.fftn
ifftn = fftpkg.ifftn
fft2 = fftpkg.fft2
ifft2 = fftpkg.ifft2

norm = linpkg.norm
inv = linpkg.inv
svd = linpkg.svd
solve = linpkg.solve
det = linpkg.det
eig = linpkg.eig
eigvals = linpkg.eigvals
eigh = linpkg.eigh
eigvalsh = linpkg.eigvalsh
lstsq = linpkg.lstsq
pinv = linpkg.pinv
cholesky = linpkg.cholesky

_restore_dict = {}

def register_func(name, func):
    if name not in __all__:
        raise ValueError, "%s not a dual function." % name
    f = sys._getframe(0).f_globals
    _restore_dict[name] = f[name]
    f[name] = func

def restore_func(name):
    if name not in __all__:
        raise ValueError, "%s not a dual function." % name
    try:
        val = _restore_dict[name]
    except KeyError:
        return
    else:
        sys._getframe(0).f_globals[name] = val

def restore_all():
    for name in _restore_dict.keys():
        restore_func(name)