Major refactoring for JAX-style classes.

essos/dynamics.py

@@ -1,3 +1,4 @@
+from pyexpat import model
 import jax
 jax.config.update("jax_enable_x64", True)
 import jax.numpy as jnp
@@ -53,7 +54,7 @@ def compute_orbit_params(xyz, vpar):
 class Particles():
     def __init__(self, initial_xyz=None, initial_vparallel_over_v=None, charge=ALPHA_PARTICLE_CHARGE,
                  mass=ALPHA_PARTICLE_MASS, energy=FUSION_ALPHA_PARTICLE_ENERGY, min_vparallel_over_v=-1,
-                 max_vparallel_over_v=1, field=None, initial_vxvyvz=None, initial_xyz_fullorbit=None):
+                 max_vparallel_over_v=1, field=None, initial_vxvyvz=None, initial_xyz_fullorbit=None, phase_angle_full_orbit = 0):
         self.charge = charge
         self.mass = mass
         self.energy = energy
@@ -85,6 +86,21 @@ def to_full_orbit(self, field):
         self.initial_xyz_fullorbit, self.initial_vxvyvz = gc_to_fullorbit(field=field, initial_xyz=self.initial_xyz, initial_vparallel=self.initial_vparallel,
                                                                             total_speed=self.total_speed, mass=self.mass, charge=self.charge,
                                                                             phase_angle_full_orbit=self.phase_angle_full_orbit)
+
+    def join(self, other, field=None):
+        assert isinstance(other, Particles), "Cannot join with non-Particles object"
+        assert self.charge == other.charge, "Cannot join particles with different charges"
+        assert self.mass == other.mass, "Cannot join particles with different masses"
+        assert self.energy == other.energy, "Cannot join particles with different energies"
+
+        charge = self.charge
+        mass = self.mass
+        energy = self.energy
+        initial_xyz = jnp.concatenate((self.initial_xyz, other.initial_xyz), axis=0)
+        initial_vparallel_over_v = jnp.concatenate((self.initial_vparallel_over_v, other.initial_vparallel_over_v), axis=0)
+
+        return Particles(initial_xyz=initial_xyz, initial_vparallel_over_v=initial_vparallel_over_v, charge=charge, mass=mass, energy=energy, field=field)
+
 @partial(jit, static_argnums=(2))
 def GuidingCenterCollisionsDiffusionMu(t,
                   initial_condition,
@@ -584,50 +600,6 @@ def condition_BioSavart(t, y, args, **kwargs):
 
         self._trajectories = self.trace()
 
-        if self.particles is not None:
-            self.energy = jnp.zeros((self.particles.nparticles, self.times_to_trace))
-
-        if model == 'GuidingCenter' or  model == 'GuidingCenterAdaptative' :
-            @jit
-            def compute_energy_gc(trajectory):
-                xyz = trajectory[:, :3]
-                vpar = trajectory[:, 3]
-                AbsB = vmap(self.field.AbsB)(xyz)
-                mu = (self.particles.energy - self.particles.mass * vpar[0]**2 / 2) / AbsB[0]
-                return self.particles.mass * vpar**2 / 2 + mu * AbsB
-            self.energy = vmap(compute_energy_gc)(self._trajectories)         
-        elif model == 'GuidingCenterCollisions':
-            @jit
-            def compute_energy_gc(trajectory):
-                return 0.5*self.particles.mass* trajectory[:, 3]**2
-            self.energy = vmap(compute_energy_gc)(self._trajectories)
-        elif model == 'GuidingCenterCollisionsMuIto' or model == 'GuidingCenterCollisionsMuFixed' or model == 'GuidingCenterCollisionsMuAdaptative' :
-            @jit
-            def compute_energy_gc(trajectory):
-                xyz = trajectory[:, :3]                
-                vpar = trajectory[:, 3]*SPEED_OF_LIGHT
-                mu = trajectory[:, 4]*self.particles.mass*SPEED_OF_LIGHT**2
-                AbsB = vmap(self.field.AbsB)(xyz)
-                return self.particles.mass * vpar**2 / 2 + mu*AbsB
-            self.energy = vmap(compute_energy_gc)(self._trajectories)
-            @jit
-            def compute_vperp_gc(trajectory):
-                xyz = trajectory[:, :3]                
-                mu = trajectory[:, 4]*self.particles.mass*SPEED_OF_LIGHT**2
-                AbsB = vmap(self.field.AbsB)(xyz)
-                return jnp.sqrt(2.*mu*AbsB/self.particles.mass)
-            self.vperp_final = vmap(compute_vperp_gc)(self._trajectories)     
-        elif model == 'FullOrbit' or model == 'FullOrbit_Boris' or model == 'FullOrbitCollisions':
-            @jit
-            def compute_energy_fo(trajectory):
-                vxvyvz = trajectory[:, 3:]
-                return self.particles.mass / 2 * (vxvyvz[:, 0]**2 + vxvyvz[:, 1]**2 + vxvyvz[:, 2]**2)
-            self.energy = vmap(compute_energy_fo)(self._trajectories)
-        elif model == 'FieldLine' or model== 'FieldLineAdaptative':
-            self.energy = jnp.ones((len(initial_conditions), self.times_to_trace))
-
-
-
         self.trajectories_xyz = vmap(lambda xyz: vmap(lambda point: self.field.to_xyz(point[:3]))(xyz))(self.trajectories)
 
         if isinstance(field, Vmec):
@@ -641,11 +613,8 @@ def compute_energy_fo(trajectory):
             else:                
                 self.loss_fractions, self.total_particles_lost, self.lost_times = self.loss_fraction_BioSavart(boundary)
         else:
-            self.loss_fractions = None
-            self.total_particles_lost = None
-            self.loss_times = None
+            self.trajectories_xyz = self.trajectories
 
-    @partial(jit, static_argnums=(0))
     def trace(self):
         @jit
         def compute_trajectory(initial_condition, particle_key) -> jnp.ndarray:
@@ -848,7 +817,7 @@ def update_state(state, _):
                     solver=diffrax.Dopri8(),
                     args=self.args,
                     saveat=SaveAt(ts=self.times),
-                    throw=False,
+                    throw=True,
                     # adjoint=DirectAdjoint(),
                     progress_meter=self.progress_meter,
                     max_steps=10000000000,
@@ -871,15 +840,41 @@ def trajectories(self):
     def trajectories(self, value):
         self._trajectories = value
 
-    def _tree_flatten(self):
-        children = (self.trajectories,)  # arrays / dynamic values
-        aux_data = {'field': self.field, 'model': self.model}  # static values
-        return (children, aux_data)
-
-    @classmethod
-    def _tree_unflatten(cls, aux_data, children):
-        return cls(*children, **aux_data)
+    def energy(self):
+        assert 'GuidingCenter' in self.model or 'FullOrbit' in self.model, "Energy calculation is only available for GuidingCenter and FullOrbit models"
+        mass = self.particles.mass
+
+        if self.model == 'GuidingCenter' or self.model == 'GuidingCenterAdaptative' or \
+           self.model == 'GuidingCenterCollisionsMuIto' or self.model == 'GuidingCenterCollisionsMuFixed' or self.model == 'GuidingCenterCollisionsMuAdaptative':
+            initial_xyz = self.initial_conditions[:, :3]
+            initial_vparallel = self.initial_conditions[:, 3]
+            initial_B = vmap(self.field.AbsB)(initial_xyz)
+            mu_array = (self.particles.energy - 0.5 * mass * jnp.square(initial_vparallel)) / initial_B
+            def compute_energy(trajectory, mu):
+                xyz = trajectory[:, :3]
+                vpar = trajectory[:, 3]
+                AbsB = vmap(self.field.AbsB)(xyz)                
+                return 0.5 * mass * jnp.square(vpar) + mu * AbsB
+
+            energy = vmap(compute_energy)(self.trajectories, mu_array)
 
+        elif self.model == 'GuidingCenterCollisions':
+            def compute_energy(trajectory):
+                return 0.5 * mass * trajectory[:, 3]**2
+            energy = vmap(compute_energy)(self.trajectories)
+
+        elif self.model == 'FullOrbit':
+            def compute_energy(trajectory):
+                vxvyvz = trajectory[:, 3:]
+                v_squared = jnp.sum(jnp.square(vxvyvz), axis=1)
+                return 0.5 * mass * v_squared
+            energy = vmap(compute_energy)(self.trajectories)
+
+        elif self.model == 'FieldLine' or self.model == 'FieldLineAdaptative':
+            energy = jnp.ones((len(self.initial_conditions), self.times_to_trace))
+
+        return energy
+
     def to_vtk(self, filename):
         try: import numpy as np
         except ImportError: raise ImportError("The 'numpy' library is required. Please install it using 'pip install numpy'.")
@@ -899,7 +894,7 @@ def plot(self, ax=None, show=True, axis_equal=True, n_trajectories_plot=5, **kwa
         trajectories_xyz = jnp.array(self.trajectories_xyz)
         n_trajectories_plot = jnp.min(jnp.array([n_trajectories_plot, trajectories_xyz.shape[0]]))
         for i in random.choice(random.PRNGKey(0), trajectories_xyz.shape[0], (n_trajectories_plot,), replace=False):
-            ax.plot(trajectories_xyz[i, :, 0], trajectories_xyz[i, :, 1], trajectories_xyz[i, :, 2], linewidth=0.5, **kwargs)
+            ax.plot(trajectories_xyz[i, :, 0], trajectories_xyz[i, :, 1], trajectories_xyz[i, :, 2], **kwargs)
         ax.grid(False)
         if axis_equal:
             fix_matplotlib_3d(ax)
@@ -966,7 +961,7 @@ def poincare_plot(self, shifts = [jnp.pi/2], orientation = 'toroidal', length =
         """
         Plot Poincare plots using scipy to find the roots of an interpolation. Can take particle trace or field lines.
         Args:
-            shifts (list, optional): Apply a linear shift to dependent data. Default is [0].
+            shifts (list, optional): Apply a linear shift to dependent data. Default is [pi/2].
             orientation (str, optional): 
                 'toroidal' - find time values when toroidal angle = shift [0, 2pi].
                 'z' - find time values where z coordinate = shift. Default is 'toroidal'.
@@ -1053,7 +1048,18 @@ def process_trajectory(X_i, Y_i, T_i):
             plt.show()
 
         return plotting_data
-
+
+    def _tree_flatten(self):
+        children = (self.trajectories, self.initial_conditions, self.times)  # arrays / dynamic values
+        aux_data = {'field': self.field, 'electric_field': self.electric_field, 'model': self.model, 'maxtime': self.maxtime, 'timestep': self.timestep,
+                    'rtol': self.rtol, 'atol': self.atol, 'particles': self.particles, 'condition': self.condition, 'tag_gc': self.tag_gc}  # static values
+        return (children, aux_data)
+
+    @classmethod
+    def _tree_unflatten(cls, aux_data, children):
+        return cls(*children, **aux_data)
+
+
 tree_util.register_pytree_node(Tracing,
                                Tracing._tree_flatten,
                                Tracing._tree_unflatten)