Merge pull request #3 from tpugnat/particle_generation

Particle generation. Thanks Thomas!

xdyna/da.py

@@ -428,28 +428,27 @@ def generate_initial_radial(self, *, angles, r_min, r_max, r_step=None, r_num=No
 
 
     # Not allowed on parallel process
-    def generate_random_initial(self, *, num_part=1000, r_max=25, px_norm=0, py_norm=0, zeta=0, delta=0.00027,
-                                normalised_emittance=None, nseeds=None, pairs_shift=0, pairs_shift_var=None):
+    def generate_random_initial(self, *, num_part=1000, r_max=25, px_norm=0, py_norm=0, zeta=0, delta=0.00027, ang_min=None,
+                                ang_max=None, normalised_emittance=None, nseeds=None, pairs_shift=0, pairs_shift_var=None):
         """Generate the initial conditions in a 2D random grid.
         """
 
         self._prepare_generation(normalised_emittance, nseeds, pairs_shift, pairs_shift_var)
 
         # Make the data
         rng = default_rng()
+        ang_min = 0 if ang_min is None else ang_min
+        ang_max = 90 if ang_max is None else ang_max
         if self.meta.nseeds > 0:
             r = rng.uniform(low=0, high=r_max**2, size=num_part*self.meta.nseeds)
-            r = np.sqrt(r)
-            th = rng.uniform(low=0, high=2*np.pi, size=num_part*self.meta.nseeds)
-            x = r*np.cos(th)
-            y = r*np.sin(th)
+            th = rng.uniform(low=ang_min, high=ang_max*np.pi/180, size=num_part*self.meta.nseeds)
             seeds = np.repeat(np.arange(1,self.meta.nseeds+1), num_part)
         else:
             r = rng.uniform(low=0, high=r_max**2, size=num_part)
-            r = np.sqrt(r)
-            th = rng.uniform(low=0, high=2*np.pi, size=num_part)
-            x = r*np.cos(th)
-            y = r*np.sin(th)
+            th = rng.uniform(low=ang_min, high=ang_max*np.pi/180, size=num_part)
+        r = np.sqrt(r)
+        x = r*np.cos(th)
+        y = r*np.sin(th)
 
         # Make dataframe
         self._surv = pd.DataFrame()