Adding plot to 2d plot python example

henryiii · henryiii · commit 2437af9958e8 · 2018-01-22T11:40:34.000+01:00
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -1,13 +1,13 @@
 cmake_minimum_required(VERSION 3.4)
 
 project(GOOFIT
-    VERSION 2.1.0
+    VERSION 2.1.1
     LANGUAGES CXX)
 
-#set(GOOFIT_TAG "dev")
+set(GOOFIT_TAG "dev")
 #set(GOOFIT_TAG "alpha")
 #set(GOOFIT_TAG "beta")
-set(GOOFIT_TAG "release")
+#set(GOOFIT_TAG "release")
 
 ### Require out-of-source builds
 file(TO_CMAKE_PATH "${PROJECT_BINARY_DIR}/CMakeLists.txt" LOC_PATH)
diff --git a/python/PDFs/GooPdf.cu b/python/PDFs/GooPdf.cu
@@ -28,5 +28,7 @@ void init_GooPdf(py::module &m) {
              R"raw(
                 Run makeGrid, set data, evaluateAtPoints, then recover original data.
                 )raw")
-        .def("setFitControl", &GooPdf::setFitControl, "Set a fit control.", "fit_control"_a);
+        .def("setFitControl", &GooPdf::setFitControl, "Set a fit control.", "fit_control"_a)
+        .def("getCompProbsAtDataPoints", &GooPdf::getCompProbsAtDataPoints, "Returns the probability at the current data points, supports multidimensional datasets.");
+    ;
 }
diff --git a/python/examples/dalitz.py b/python/examples/dalitz.py
@@ -26,6 +26,7 @@ def cpuGetM23(massPZ,massPM):
 
 
 def getToyData(toyFileName):
+    print("Reading {} ...".format(toyFileName))
     out = pd.read_csv(toyFileName,
             delim_whitespace=True,
             comment='=',
@@ -234,15 +235,28 @@ def runToyFit(toyFileName):
     signal = makeSignalPdf(m12, m13, eventNumber)
     signal.setData(data)
     signal.setDataSize(data.getNumEvents())
-    datapdf = FitManager(signal)
-
-    datapdf.fit()
-    return datapdf
+    fitman = FitManager(signal)
 
+    fitman.fit()
+    return fitman, signal, m12, m13
 
 def main():
-    print("main")
     filename = GOOFIT_SOURCE_DIR + "/examples/dalitz/dalitz_toyMC_000.txt"
-    return runToyFit(filename)
-
-sys.exit(int(int(main())))
+    fitman, signal, m12, m23 =  runToyFit(filename)
+    
+    # Make a grid and evaluate over it
+    grid = signal.makeGrid()
+    signal.setData(grid)
+    # Segfault here:
+    #val = signal.getCompProbsAtDataPoints()
+    
+    
+    #for i in range(len(grid)):
+    #    grid.loadEvent(i)
+    #    print(m12, m13, val[0][i])
+    
+    # Double int's needed here since int may return a long and sys.exit needs a real int for Python2
+    return int(int(fitman))
+
+if __name__ == '__main__':
+    sys.exit(main())
diff --git a/python/examples/plot_2d.py b/python/examples/plot_2d.py
@@ -6,6 +6,10 @@
 from goofit import *
 import numpy as np
 
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+
 print_goofit_info()
 
 # Create random data
@@ -44,7 +48,25 @@
 print(ymean)
 print(ysigm)
 
-print(data.to_numpy())
 grid = total.makeGrid()
-print(grid.to_numpy())
 total.setData(grid)
+
+val = total.getCompProbsAtDataPoints()
+
+mat = np.array(val[0]).reshape([xvar.numbins, yvar.numbins])
+
+fig, axs = plt.subplots(1, 2, figsize=(10,5))
+
+axs[0].hist2d(xyarr[0], xyarr[1], bins=(100,100), range=((xvar.lowerlimit, xvar.upperlimit), (yvar.lowerlimit, yvar.upperlimit)))
+axs[1].imshow(mat, extent=(xvar.lowerlimit, xvar.upperlimit, yvar.lowerlimit, yvar.upperlimit), origin='lower')
+
+for ax in axs:
+    ax.set_xlabel("x")
+    ax.set_ylabel("y")
+    ax.axis("equal")
+
+axs[0].set_title("Original input (histogram)")
+axs[1].set_title("Fitted function")
+
+fig.savefig("plot_2d.png")
+
diff --git a/python/goofit/UnbinnedDataSet.cpp b/python/goofit/UnbinnedDataSet.cpp
@@ -64,5 +64,7 @@ void init_UnbinnedDataSet(py::module &m) {
              &UnbinnedDataSet::from_matrix<Eigen::Matrix<fptype, Eigen::Dynamic, Eigen::Dynamic>>,
              "Append a matrix to a dataset. The final parameter will be a count if the matrix is missing one row."
              "matrix"_a,
-             "filter"_a = false);
+             "filter"_a = false)
+        .def("loadEvent", &UnbinnedDataSet::loadEvent, "Load an event into the observables", "event_number"_a)
+    ;
 }

Original file line number	Diff line number	Diff line change
`@@ -28,5 +28,7 @@ void init_GooPdf(py::module &m) {`
`28`	`28`	`R"raw(`
`29`	`29`	`Run makeGrid, set data, evaluateAtPoints, then recover original data.`
`30`	`30`	`)raw")`
`31`		`- .def("setFitControl", &GooPdf::setFitControl, "Set a fit control.", "fit_control"_a);`
	`31`	`+ .def("setFitControl", &GooPdf::setFitControl, "Set a fit control.", "fit_control"_a)`
	`32`	`+ .def("getCompProbsAtDataPoints", &GooPdf::getCompProbsAtDataPoints, "Returns the probability at the current data points, supports multidimensional datasets.");`
	`33`	`+ ;`
`32`	`34`	`}`
Original file line number	Diff line number	Diff line change
`@@ -64,5 +64,7 @@ void init_UnbinnedDataSet(py::module &m) {`
`64`	`64`	`&UnbinnedDataSet::from_matrix<Eigen::Matrix<fptype, Eigen::Dynamic, Eigen::Dynamic>>,`
`65`	`65`	`"Append a matrix to a dataset. The final parameter will be a count if the matrix is missing one row."`
`66`	`66`	`"matrix"_a,`
`67`		`- "filter"_a = false);`
	`67`	`+ "filter"_a = false)`
	`68`	`+ .def("loadEvent", &UnbinnedDataSet::loadEvent, "Load an event into the observables", "event_number"_a)`
	`69`	`+ ;`
`68`	`70`	`}`