add mean and var functions to ParticleAnalyzer

postpic/analyzer/particles.py

@@ -557,6 +557,26 @@ def angle_xaxis(self):
     angle_xaxis.unit = 'rad'
     angle_xaxis.name = 'angle_xaxis'
 
+    # ---- Functions for measuring particle collection related values
+
+    def mean(self, func, weights=1.0):
+        '''
+        the mean of a value given by the function func. The particle weight
+        of the individual particles will be included in the calculation.
+        An additional weight can be given as well.
+        '''
+        w = self.weight() * weights
+        return np.average(func(self), weights=w)
+
+    def var(self, func, weights=1.0):
+        '''
+        variance
+        '''
+        w = self.weight() * weights
+        data = func(self)
+        m = np.average(data, weights=w)
+        return np.average((data - m)**2, weights=w)
+
     # ---- Functions to create a Histogram. ---
 
     def createHistgram1d(self, scalarfx, optargsh={'bins': 300},

test/test_particles.py

@@ -18,5 +18,14 @@ def test_dummyreader(self):
     def test_pa(self):
         self.assertEqual(len(self.pa), 10000)
 
+    def test_mean(self):
+        self.assertAlmostEqual(self.pa.mean(pa.ParticleAnalyzer.X), -0.0184337)
+        self.assertAlmostEqual(self.pa.mean(pa.ParticleAnalyzer.X, weights=self.pa.beta()), -0.01965867)
+
+    def test_var(self):
+        self.assertAlmostEqual(self.pa.var(pa.ParticleAnalyzer.X),0.97526797)
+        self.assertAlmostEqual(self.pa.var(pa.ParticleAnalyzer.Y),0.98615759)
+        self.assertRaises(KeyError, self.pa.var, pa.ParticleAnalyzer.Z)
+
 if __name__ == '__main__':
     unittest.main()