PINNs/test_bsplines.py at main · ion-g-ion/PINNs · 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
import numpy as np
import jax
import jax.numpy as jnp
import pinns
import matplotlib.pyplot as plt
import unittest

import jax.config
jax.config.update("jax_enable_x64", True)


if __name__ == '__main__':

    basis = pinns.functions.BSplineBasisJAX(np.linspace(0,1,7),2)

    x = np.linspace(0,1,1001)

    # y = basis(x)
    # plt.figure()
    # plt.plot(x,y.T)

    # y = basis(x, derivative = True)
    # plt.figure()
    # plt.plot(x,y.T)
    #
    #
    # basis = pinns.functions.BSplineBasisJAX(np.array([0,0.1,0.5,0.5,0.9,1]),2)
    #
    # y = basis(x)
    # plt.figure()
    # plt.plot(x,y.T)

    # y = basis(x, derivative = True)
    # plt.figure()
    # plt.plot(x,y.T)
    #
    #
    #
    #
    # # JAX Bspl

    # basis = pinns.functions.BSplineBasisJAX(np.linspace(0,1,7),2)
    #
    # basisj = jax.jit(basis)
    #
    # x = np.linspace(0,1,1001)

    # y = np.array(basisj(x))
    # plt.figure()
    # plt.plot(x,y.T)

    #
    # y = basis(x, derivative = True)
    # plt.figure()
    # plt.plot(x,y.T)
    #
    #
    # basis = pinns.functions.BSplineBasisJAX(np.array([0,0.1,0.5,0.5,0.9,1]),2)
    #
    # y = basis(x)
    # print(y.shape)
    #
    # plt.figure()
    # plt.plot(x,y.T)

    # y = basis(x, derivative = True)
    # plt.figure()
    # plt.plot(x,y.T)
    #
    # plt.show()


    basis = pinns.functions.PiecewiseBernsteinBasisJAX(np.array([0,0.1,0.3,0.7,0.9,1]),3)

    y = basis(x)
    print(y.shape)

    plt.figure()
    plt.title("Bernstein")
    plt.plot(x,y.T)

    y = basis(x, derivative = True)
    plt.figure()
    plt.plot(x,y.T)

    plt.show()