xref: /petsc/src/binding/petsc4py/demo/legacy/taosolve/chwirut.py (revision 5a48edb989d3ea10d6aff6c0e26d581c18691deb)
1*55a74a43SLisandro Dalcinimport sys, petsc4py
2*55a74a43SLisandro Dalcinpetsc4py.init(sys.argv)
3*55a74a43SLisandro Dalcin
4*55a74a43SLisandro Dalcinimport numpy as np
5*55a74a43SLisandro Dalcinfrom petsc4py import PETSc
6*55a74a43SLisandro Dalcin
7*55a74a43SLisandro Dalcinclass Chwirut(object):
8*55a74a43SLisandro Dalcin
9*55a74a43SLisandro Dalcin    """
10*55a74a43SLisandro Dalcin    Finds the nonlinear least-squares solution to the model
11*55a74a43SLisandro Dalcin    y = exp(-b1*x)/(b2+b3*x)  +  e
12*55a74a43SLisandro Dalcin    """
13*55a74a43SLisandro Dalcin
14*55a74a43SLisandro Dalcin    def __init__(self):
15*55a74a43SLisandro Dalcin        BETA = [0.2, 0.12, 0.08]
16*55a74a43SLisandro Dalcin        NOBSERVATIONS = 100
17*55a74a43SLisandro Dalcin        NPARAMETERS = 3
18*55a74a43SLisandro Dalcin
19*55a74a43SLisandro Dalcin        np.random.seed(456)
20*55a74a43SLisandro Dalcin        x = np.random.rand(NOBSERVATIONS)
21*55a74a43SLisandro Dalcin        e = np.random.rand(NOBSERVATIONS)
22*55a74a43SLisandro Dalcin
23*55a74a43SLisandro Dalcin        y = np.exp(-BETA[0]*x)/(BETA[1] + BETA[2]*x) + e
24*55a74a43SLisandro Dalcin
25*55a74a43SLisandro Dalcin        self.NOBSERVATIONS = NOBSERVATIONS
26*55a74a43SLisandro Dalcin        self.NPARAMETERS = NPARAMETERS
27*55a74a43SLisandro Dalcin        self.x = x
28*55a74a43SLisandro Dalcin        self.y = y
29*55a74a43SLisandro Dalcin
30*55a74a43SLisandro Dalcin    def createVecs(self):
31*55a74a43SLisandro Dalcin        X = PETSc.Vec().create(PETSc.COMM_SELF)
32*55a74a43SLisandro Dalcin        X.setSizes(self.NPARAMETERS)
33*55a74a43SLisandro Dalcin        F = PETSc.Vec().create(PETSc.COMM_SELF)
34*55a74a43SLisandro Dalcin        F.setSizes(self.NOBSERVATIONS)
35*55a74a43SLisandro Dalcin        return X, F
36*55a74a43SLisandro Dalcin
37*55a74a43SLisandro Dalcin    def formInitialGuess(self, X):
38*55a74a43SLisandro Dalcin        X[0] = 0.15
39*55a74a43SLisandro Dalcin        X[1] = 0.08
40*55a74a43SLisandro Dalcin        X[2] = 0.05
41*55a74a43SLisandro Dalcin
42*55a74a43SLisandro Dalcin    def formResidual(self, tao, X, F):
43*55a74a43SLisandro Dalcin        x, y = self.x, self.y
44*55a74a43SLisandro Dalcin        b1, b2, b3 = X.array
45*55a74a43SLisandro Dalcin        F.array = y - np.exp(-b1*x)/(b2 + b3*x)
46*55a74a43SLisandro Dalcin
47*55a74a43SLisandro Dalcin    def plotSolution(self, X):
48*55a74a43SLisandro Dalcin        try:
49*55a74a43SLisandro Dalcin            from matplotlib import pylab
50*55a74a43SLisandro Dalcin        except ImportError:
51*55a74a43SLisandro Dalcin            return
52*55a74a43SLisandro Dalcin        b1, b2, b3 = X.array
53*55a74a43SLisandro Dalcin        x, y = self.x, self.y
54*55a74a43SLisandro Dalcin        u = np.linspace(x.min(), x.max(), 100)
55*55a74a43SLisandro Dalcin        v = np.exp(-b1*u)/(b2+b3*u)
56*55a74a43SLisandro Dalcin        pylab.plot(x, y, 'ro')
57*55a74a43SLisandro Dalcin        pylab.plot(u, v, 'b-')
58*55a74a43SLisandro Dalcin        pylab.show()
59*55a74a43SLisandro Dalcin
60*55a74a43SLisandro DalcinOptDB = PETSc.Options()
61*55a74a43SLisandro Dalcin
62*55a74a43SLisandro Dalcinuser = Chwirut()
63*55a74a43SLisandro Dalcin
64*55a74a43SLisandro Dalcinx, f = user.createVecs()
65*55a74a43SLisandro Dalcinx.setFromOptions()
66*55a74a43SLisandro Dalcinf.setFromOptions()
67*55a74a43SLisandro Dalcin
68*55a74a43SLisandro Dalcintao = PETSc.TAO().create(PETSc.COMM_SELF)
69*55a74a43SLisandro Dalcintao.setType(PETSc.TAO.Type.POUNDERS)
70*55a74a43SLisandro Dalcintao.setResidual(user.formResidual, f)
71*55a74a43SLisandro Dalcintao.setFromOptions()
72*55a74a43SLisandro Dalcin
73*55a74a43SLisandro Dalcinuser.formInitialGuess(x)
74*55a74a43SLisandro Dalcintao.solve(x)
75*55a74a43SLisandro Dalcin
76*55a74a43SLisandro Dalcinplot = OptDB.getBool('plot', False)
77*55a74a43SLisandro Dalcinif plot: user.plotSolution(x)
78*55a74a43SLisandro Dalcin
79*55a74a43SLisandro Dalcinx.destroy()
80*55a74a43SLisandro Dalcinf.destroy()
81*55a74a43SLisandro Dalcintao.destroy()
82