pymodules/FEM.py at master · yaoyuki/pymodules · GitHub

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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

# coding: utf-8

# In[1]:

import numpy as np
import numba as nb

#my library
from orderedTableSearch import locate_grid, locate_on_grids, locate


# In[2]:

@nb.jit(nopython = True)
def fem_elmeval(x, x0, x1,para0, para1):
    diff = x1 - x0
    bs0 = (x1 - x)/diff
    bs1 = (x - x0)/diff
    #if bs0 >= 0.0 and bs1 >= 0.0:
    return para0*bs0 + para1*bs1
    #else:
    #    return 0.0

@nb.jit(nopython=True)
def fem_peval(xval, nodes, para):
    ie = locate(xval, nodes)
    return fem_elmeval(xval, nodes[ie], nodes[ie+1], para[ie], para[ie+1])

@nb.jit(nopython = True)
def fem_grideval(xvals, nodes, para):
    #I should prohibit extrapolation

    N = len(nodes)
    M = len(xvals)

    if N != len(para):
        print('error: N != len(para)')
        return np.ones(M)*np.nan
#         return None #this sometimes causes a problem.


    #xelements = locate_grid(xvals, nodes)
    xelements = locate_on_grids(xvals, nodes)
    ans = np.zeros(M)
    for ix,x in enumerate(xvals):
        ie = xelements[ix]

        ans[ix] = fem_elmeval(x, nodes[ie], nodes[ie+1], para[ie], para[ie+1])

    return ans

#@nb.generated_jit(nopython=True)
#def femeval(xvals, nodes, para):
#    #I should prohibit extrapolation
#    if not hasattr(xvals, "__len__"): #if xvals is scalar
#        ie = locate_grid(xvals, nodes)
#
#        return fem_elmeval(xvals, nodes[ie], nodes[ie+1], para[ie], para[ie+1])
#    else: #if xvals is actually array-like
#
#        return fem_grideval(xvals, nodes, para)

# @nb.generated_jit(nopython=True)
@nb.jit(nopython=True)
def femeval(xvals, nodes, para):
    #I should prohibit extrapolation
    # if isinstance(xvals, nb.types.Float) or isinstance(xvals, nb.types.Integer):
    if isinstance(xvals, float) or isinstance(xvals, int):

        return fem_peval(xvals, nodes, para)

    else: #if xvals is actually array-like

        return fem_grideval(xvals, nodes, para)

@nb.jit(nopython = True)
def dy_femeval(x, xnodes):
    """
    derivative wrt parameters (y)
    or returns only the weights on xnodes

    input:
    x: a point to be interpolated
    xnode: the grid

    output:
    ans (N-dim array)

    """

    N = len(xnodes)


    ans = np.zeros(N)

    ie = locate(x, xnodes) #need to think of using hunt

    #need to add some check
    x0 = xnodes[ie]
    x1 = xnodes[ie+1]
    diff = x1 - x0

    ans[ie] = (x1 - x)/diff
    ans[ie+1] = (x - x0)/diff

    return ans