Merge pull request #24 from PyCOMPLETE/develop

New interface TransverseMap
PyCOMPLETE · Jan 11, 2016 · cfbfe83 · cfbfe83
2 parents ad404d1 + cac79e1
commit cfbfe83
Show file tree

Hide file tree

Showing 26 changed files with 712 additions and 663 deletions.
diff --git a/field_maps/Transverse_Efield_map.py b/field_maps/Transverse_Efield_map.py
@@ -9,8 +9,15 @@
 class Transverse_Efield_map(object):
 	def __init__(self, xg, yg, Ex, Ey, n_slices, z_cut,
 		L_interaction, flag_clean_slices = False, wrt_slice_centroid = False):
-
-		self.slicer = UniformBinSlicer(n_slices = n_slices, z_cuts=(-z_cut, z_cut))
+
+		if type(z_cut) is float:
+			z_cuts = (-z_cut, z_cut)
+		elif type(z_cut) is tuple:
+			z_cuts = z_cut
+		else:
+			raise ValueError('Type not recognized!')
+
+		self.slicer = UniformBinSlicer(n_slices = n_slices, z_cuts=z_cuts)
 		self.L_interaction = L_interaction
 		self.flag_clean_slices = flag_clean_slices
 		self.wrt_slice_centroid = wrt_slice_centroid

diff --git a/particles/particles.py b/particles/particles.py
@@ -194,6 +194,41 @@ def reorder(self, permutation, except_for_attrs=[]):
                 continue
             reordered = getattr(self, attr)[permutation]
             setattr(self, attr, reordered)
+
+    def __add__(self, other):
+        '''Merges two beams.
+		'''
+        #print 'Checks still to be added!!!!!!'
+
+        self_coords_n_momenta_dict = self.get_coords_n_momenta_dict()
+        other_coords_n_momenta_dict = other.get_coords_n_momenta_dict()
+
+        result = Particles(macroparticlenumber=self.macroparticlenumber+other.macroparticlenumber, 
+                    particlenumber_per_mp=self.particlenumber_per_mp, charge=self.charge,
+					mass=self.mass, circumference=self.circumference, gamma=self.gamma, coords_n_momenta_dict={})
+
+
+        for coord in self_coords_n_momenta_dict.keys():
+            #setattr(result, coord, np.concatenate((self_coords_n_momenta_dict[coord].copy(), other_coords_n_momenta_dict[coord].copy())))
+            result.update({coord: np.concatenate((self_coords_n_momenta_dict[coord].copy(), other_coords_n_momenta_dict[coord].copy()))})
+
+        return result 
+
+    def __radd__(self, other):
+        if other==0:
+            self_coords_n_momenta_dict = self.get_coords_n_momenta_dict()
+            result = Particles(macroparticlenumber=self.macroparticlenumber, 
+                    particlenumber_per_mp=self.particlenumber_per_mp, charge=self.charge,
+                    mass=self.mass, circumference=self.circumference, gamma=self.gamma, coords_n_momenta_dict={})
+
+            for coord in self_coords_n_momenta_dict.keys():
+                #setattr(result, coord, np.concatenate((self_coords_n_momenta_dict[coord].copy(), other_coords_n_momenta_dict[coord].copy())))
+                result.update({coord: self_coords_n_momenta_dict[coord].copy()})
+        else:
+            result = self.__add__(other)
+
+        return result      
+
 
     # Statistics methods
 

diff --git a/testing/interactive-tests/ApertureNLossesTest.ipynb b/testing/interactive-tests/ApertureNLossesTest.ipynb
diff --git a/testing/interactive-tests/BeamOptics.ipynb b/testing/interactive-tests/BeamOptics.ipynb
diff --git a/testing/interactive-tests/DetunersCythonTest.ipynb b/testing/interactive-tests/DetunersCythonTest.ipynb
diff --git a/testing/interactive-tests/DetunersTest.ipynb b/testing/interactive-tests/DetunersTest.ipynb
diff --git a/testing/interactive-tests/GeneratorTest.ipynb b/testing/interactive-tests/GeneratorTest.ipynb
diff --git a/testing/interactive-tests/MonitorTest.ipynb b/testing/interactive-tests/MonitorTest.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 7,
    "metadata": {
     "collapsed": false
    },
@@ -15,21 +15,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 8,
    "metadata": {
     "collapsed": false
    },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "PyHEADTAIL v1.2.0-153-gbc0ea69a06-dirty\n",
-      "\n",
-      "\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "import numpy as np\n",
     "from scipy.constants import m_p, c, e\n",
@@ -45,7 +35,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 9,
    "metadata": {
     "collapsed": false
    },
@@ -143,7 +133,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 10,
    "metadata": {
     "collapsed": false
    },
@@ -177,7 +167,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 11,
    "metadata": {
     "collapsed": false
    },
@@ -197,7 +187,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 12,
    "metadata": {
     "collapsed": false
    },
@@ -206,18 +196,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "0.0598522977141\n"
-     ]
-    },
-    {
-     "ename": "KeyboardInterrupt",
-     "evalue": "",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
-      "\u001b[0;32m<ipython-input-10-39c3aff5b0d1>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     30\u001b[0m     \u001b[0mslice_monitor\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdump\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbunch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     31\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 32\u001b[0;31m     \u001b[0mslice_set_pmon\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbunch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_slices\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0munibin_slicer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     33\u001b[0m     \u001b[0marrays_dict\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mupdate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m{\u001b[0m\u001b[0;34m'slidx'\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mslice_set_pmon\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mslice_index_of_particle\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'zz'\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mbunch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mz\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     34\u001b[0m     \u001b[0mparticle_monitor\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdump\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbunch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0marrays_dict\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
+      "0.0614682430577\n"
      ]
     }
    ],
@@ -228,7 +207,7 @@
     "    alpha_0, Q_s, R)\n",
     "\n",
     "trans_map = TransverseMap(\n",
-    "    C, s, alpha_x, beta_x, D_x, alpha_y, beta_y, D_y, Q_x, Q_y)\n",
+    "    s, alpha_x, beta_x, D_x, alpha_y, beta_y, D_y, Q_x, Q_y)\n",
     "\n",
     "# Slicer config for SliceMonitor.\n",
     "unibin_slicer = UniformBinSlicer(n_slices=10, n_sigma_z=None, z_cuts=None)\n",
@@ -263,7 +242,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 13,
    "metadata": {
     "collapsed": false
    },
@@ -283,6 +262,15 @@
    },
    "outputs": [],
    "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": true
+   },
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -301,7 +289,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython2",
-   "version": "2.7.10"
+   "version": "2.7.11"
   }
  },
  "nbformat": 4,

diff --git a/testing/interactive-tests/RFBucketHamiltonianDeformation.ipynb b/testing/interactive-tests/RFBucketHamiltonianDeformation.ipynb
diff --git a/testing/interactive-tests/RFQTest.ipynb b/testing/interactive-tests/RFQTest.ipynb
diff --git a/testing/interactive-tests/SlicingTest.ipynb b/testing/interactive-tests/SlicingTest.ipynb