xref: /petsc/src/binding/petsc4py/demo/regressor/test_regressor.py (revision c12c126234ed623246a63bfa78c9f75a3aa00323)
1# Simple linear regression test
2# =============================
3#
4# Use the petsc4py interface to PetscRegressor to solve the same quadratic
5# polynomial data-fitting problem that is solved in src/ml/regressor/tests/ex1.c.
6
7import sys
8import petsc4py
9
10petsc4py.init(sys.argv)
11
12from petsc4py import PETSc
13
14comm = PETSc.COMM_WORLD
15size = comm.getSize()
16rank = comm.getRank()
17
18OptDB = PETSc.Options()
19
20y_array = [0.2, -0.3, -0.8, -0.3, 1.2]
21X_array = [-1.0, 0.5, -0.5, -0.25, 0.0, -0.5, 0.5, -0.25, 1.0, 0.5]
22rows_ix = [0, 1, 2, 3, 4]
23cols_ix = [0, 1]
24
25X = PETSc.Mat().create(comm=comm)
26X.setSizes((5,2))
27X.setFromOptions()
28X.setUp()
29
30y = PETSc.Vec().create(comm=comm)
31y.setSizes(5)
32y.setFromOptions()
33
34if not rank :
35    X.setValues(rows_ix,cols_ix,X_array,addv=True)
36    y.setValues(rows_ix,y_array,addv=False)
37
38X.assemblyBegin(X.AssemblyType.FINAL)
39X.assemblyEnd(X.AssemblyType.FINAL)
40y.assemblyBegin()
41y.assemblyEnd()
42
43y_predicted = y.duplicate()
44
45regressor = PETSc.Regressor().create(comm=comm)
46regressor.setType(PETSc.Regressor.Type.LINEAR)
47regressor.fit(X,y)
48regressor.predict(X,y_predicted)
49y_predicted.view()
50
51
52
53