Slim down all and make it more Pythonic. Easy to use now with example for octupoles.

Author: like2000

Detuning.py

@@ -1,17 +1,17 @@
+from __future__ import division
+
 import numpy as np
 
 
-class Detuning:
-    def __call__(self,jx,jy):
-        raise NotImplemented;
+class Detuning(object):
+
+    def __call__(self, jx, jy):
+        raise NotImplemented
 
 
 class LinearDetuning(Detuning):
-    _startTune = 0.31;
-    _slopex = 0.0;
-    _slopey = 0.0;
 
-    def __init__(self, startTune, slopex, slopey):
+    def __init__(self, startTune=0.31, slopex=0, slopey=0):
         self._startTune = startTune
         self._slopex = slopex
         self._slopey = slopey
@@ -21,11 +21,8 @@ def __call__(self, jx, jy):
 
 
 class FootprintDetuning(Detuning):
-    _footprint = None;
-    _plane = None;
-
     # 0 for H and 1 for V
-    def __init__(self, footprint, plane=0):
+    def __init__(self, footprint=None, plane=0):
         self._footprint = footprint
         self._plane = plane
 

Dispersion.py

@@ -1,23 +1,21 @@
-class Dispersion:
-    _detuning = None;
-    _distribution = None;
-    _Q = 0.0;
-    _epsilon = 1E-6;
+from __future__ import division
 
-    def __init__(self, distribution, detuning, Q, epsilon=1E-6):
+
+class Dispersion(object):
+
+    def __init__(self, distribution=None, detuning=None, Q=0, epsilon=1E-6):
         self._distribution = distribution
         self._detuning = detuning
         self._Q = Q
         self._epsilon = epsilon
 
-    def setEpsilon(self,epsilon):
+    def setEpsilon(self, epsilon):
         self._epsilon = epsilon
 
     def getEpsilon(self):
         return self._epsilon
 
     def getValue(self, jx, jy):
-
         # DI = (
         #      (jx * self._distribution.getDJx(jx,jy)) /
         #      (complex(self._Q - self._detuning(jx,jy), self._epsilon))

Distribution.py

@@ -1,4 +1,5 @@
 import numpy as np
+
 from abc import ABCMeta, abstractmethod
 
 
@@ -25,10 +26,10 @@ class BiGaussian(Distribution):
 
     def getValue(self, jx, jy):
         '''Return pdf of normal bi-Gaussian'''
-        return np.exp(-(jx + jy));
+        return np.exp(-(jx + jy))
 
     def getDJx(self, jx, jy):
-        return -self.getValue(jx, jy);
+        return -self.getValue(jx, jy)
 
     def getDJy(self, jx, jy):
-        return -self.getValue(jx, jy);
+        return -self.getValue(jx, jy)

Integrator.py

@@ -1,5 +1,6 @@
 from __future__ import division
 
+import warnings
 import numpy as np
 from scipy.integrate import dblquad, simps
 
@@ -25,13 +26,13 @@ def integrate(self): pass
 
 
 class Boundary:
-    _boundary = 0.0;
+    _boundary = 0.0
 
-    def __init__(self,boundary):
-        self._boundary = boundary;
+    def __init__(self, boundary):
+        self._boundary = boundary
 
-    def __call__(self,x):
-        return self._boundary;
+    def __call__(self, x):
+        return self._boundary
 
 
 class DblquadIntegrator(Integrator):
@@ -50,10 +51,12 @@ def integrate(self, Q):
         rDispersion = RealDispersion(self.distribution, self.detuning, Q)
         iDispersion = ImaginaryDispersion(self.distribution, self.detuning, Q)
 
-        realPart, realErr = dblquad(rDispersion, self.minJx, self.maxJx,
-                                                 self.minJy, self.maxJy)
-        imagPart, imagErr = dblquad(iDispersion, self.minJx, self.maxJx,
-                                                 self.minJy, self.maxJy)
+        realPart, realErr = dblquad(rDispersion,
+                                    self.minJx, self.maxJx,
+                                    self.minJy, self.maxJy)
+        imagPart, imagErr = dblquad(iDispersion,
+                                    self.minJx, self.maxJx,
+                                    self.minJy, self.maxJy)
 
         return -1.0/complex(realPart, imagPart)
 
@@ -64,17 +67,23 @@ def __init__(self, *args, **kwargs):
 
         super(SimpsonIntegrator, self).__init__(*args, **kwargs)
 
-        if 'n_steps' not in kwargs: kwargs['n_steps'] = 2000
-        self.jx = np.linspace(self.minJx, self.maxJx, kwargs['n_steps'])
-        self.jy = np.linspace(self.minJy, self.maxJy, kwargs['n_steps'])
+        if 'n_steps' not in kwargs:
+            kwargs['n_steps'] = 1000
+        n_steps = kwargs['n_steps']
+        self.jx = np.linspace(self.minJx, self.maxJx, n_steps)
+        self.jy = np.linspace(self.minJy, self.maxJy, n_steps)
         self.JX, self.JY = np.meshgrid(self.jx, self.jy)
 
     def integrate(self, Q, epsilon=1e-6):
 
-        # dd = Dispersion(self.distribution, self.detuning, Q, epsilon=epsilon).getValue(self.JX, self.JY)
-        dd = np.array([[
-            Dispersion(self.distribution, self.detuning, Q, epsilon=epsilon).getValue(x, y)
-            for y in self.jy] for x in self.jx])
+        dd = Dispersion(
+            self.distribution, self.detuning, Q, epsilon=epsilon
+        ).getValue(self.JX, self.JY)
+        # dd = np.array([[
+        #     Dispersion(
+        #         self.distribution, self.detuning, Q, epsilon=epsilon
+        #     ).getValue(x, y)
+        #     for y in self.jy] for x in self.jx])
 
         return -1./simps(simps(dd, self.jx), self.jy)
 
@@ -85,20 +94,33 @@ def __init__(self, *args, **kwargs):
 
         super(TrapzIntegrator, self).__init__(*args, **kwargs)
 
-        if 'n_steps' not in kwargs: kwargs['n_steps'] = 1000
-        self.jx = np.linspace(self.minJx, self.maxJx, kwargs['n_steps'])
-        self.jy = np.linspace(self.minJy, self.maxJy, kwargs['n_steps'])
+        if 'n_steps' not in kwargs:
+            kwargs['n_steps'] = 1000
+        n_steps = 1000
+        self.jx = np.linspace(self.minJx, self.maxJx, n_steps)
+        self.jy = np.linspace(self.minJy, self.maxJy, n_steps)
         self.JX, self.JY = np.meshgrid(self.jx, self.jy)
 
     def integrate(self, Q, epsilon=1e-6):
 
-        dd = Dispersion(self.distribution, self.detuning, Q, epsilon=epsilon).getValue(self.JX, self.JY)
+        dd = Dispersion(
+            self.distribution, self.detuning, Q, epsilon=epsilon
+        ).getValue(self.JX, self.JY)
+
         return -1./np.trapz(np.trapz(dd, self.jx), self.jy)
 
 
 class FixedTrapezoidalIntegrator(Integrator):
 
-    def __init__(self, distribution=None, detuning=None, minJ=0, maxJ=18, nStep=2000):
+    def __init__(self, distribution=None, detuning=None,
+                 minJ=0, maxJ=18, nStep=2000):
+        warnings.simplefilter('always', DeprecationWarning)
+        warnings.warn('Class "{:s}" '.format(self.__name__) +
+                      'is deprecated and will be replaced in the ' +
+                      'near future.',
+                      category=DeprecationWarning, stacklevel=2)
+        warnings.simplefilter('default', DeprecationWarning)
+
         self._distribution = distribution
         self._detuning = detuning
         self._minJx = minJ
@@ -158,7 +180,14 @@ class AdaptiveRectangularIntegrator(Integrator):
     _arrayY = None;
 
 
-    def __init__(self,distribution,detuning,minJ=0.0,maxJ=18.0):
+    def __init__(self, distribution, detuning, minJ=0.0, maxJ=18.0):
+        warnings.simplefilter('always', DeprecationWarning)
+        warnings.warn('Class "{:s}" '.format(self.__name__) +
+                      'is deprecated and will be replaced in the ' +
+                      'near future.',
+                      category=DeprecationWarning, stacklevel=2)
+        warnings.simplefilter('default', DeprecationWarning)
+
         self._distribution = distribution;
         self._detuning = detuning;
         self._minJx = minJ;

main.py

@@ -3,11 +3,12 @@
 import sys, time
 import numpy as np
 import matplotlib.pyplot as plt
+plt.switch_backend('TkAgg')
 
 from PySSD.Distribution import BiGaussian
 from PySSD.Detuning import LinearDetuning
 from PySSD.Dispersion import Dispersion
-from PySSD.Integrator import FixedTrapezoidalIntegrator, SimpsonIntegrator
+from PySSD.Integrator import FixedTrapezoidalIntegrator, SimpsonIntegrator, TrapzIntegrator
 
 # from PySSD.Detunign import FootprintDetuning
 # from PySSD.Footprint import Footprint, parseDynapTune
@@ -62,8 +63,9 @@ def findQs(detuning, stepSize=5E-5, maxJ=18.0, dJ=0.1, margin=1):
 maxJ = 18
 Qs = findQs(detuning, stepSize=2e-5, maxJ=maxJ)
 
-# integrator = FixedTrapezoidalIntegrator(distribution, detuning, maxJ=maxJ)
+integrator = FixedTrapezoidalIntegrator(distribution, detuning, maxJ=maxJ)
 integrator = SimpsonIntegrator(distribution, detuning, maxJ=maxJ)
+integrator = TrapzIntegrator(distribution, detuning, maxJ=maxJ)
 
 tuneShifts = []
 outputFileName = './Test.sdiag';