Merge pull request #16 from uwplasma/development

Add Poincaré plot and sharding

Author: rogeriojorge

README.md

@@ -90,7 +90,7 @@ ESSOS/
 │       ├── ESSOS_bio_savart_LandremanPaulQA.json
 │       ├── SIMSOPT_bio_savart_LandremanPaulQA.json
 │       ├── wout_n3are_R7.75B5.7.nc
-│       └── wout_LandremanPaul2021_QA_reactorScale_lowres_reference.nc
+│       └── wout_LandremanPaul2021_QA_reactorScale_lowres.nc
 ├── tests/
 │   ├── test_coils.py
 │   ├── test_constants.py

essos/__main__.py

@@ -8,7 +8,8 @@
 from essos.coils import Coils, CreateEquallySpacedCurves
 from essos.fields import near_axis, BiotSavart
 from essos.dynamics import Tracing
-from essos.optimization import optimize_coils_for_nearaxis
+from essos.optimization import optimize_loss_function
+from essos.objective_functions import loss_coils_for_nearaxis
 
 
 def main(cl_args=sys.argv[1:]):
@@ -53,9 +54,10 @@ def main(cl_args=sys.argv[1:]):
     # Optimize coils
     print(f'Optimizing coils with {maximum_function_evaluations} function evaluations.')
     time0 = time()
-    coils_optimized = optimize_coils_for_nearaxis(field, coils_initial, maximum_function_evaluations=maximum_function_evaluations,
-                                                        max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature,
-                                                        tolerance_optimization=tolerance_optimization)
+    coils_optimized = optimize_loss_function(loss_coils_for_nearaxis, initial_dofs=coils_initial.x,
+                                    coils=coils_initial, tolerance_optimization=tolerance_optimization,
+                                    maximum_function_evaluations=maximum_function_evaluations, field_nearaxis=field,
+                                    max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature,)
     print(f"Optimization took {time()-time0:.2f} seconds")
 
     # Trace fieldlines
@@ -81,10 +83,10 @@ def main(cl_args=sys.argv[1:]):
     ax1 = fig.add_subplot(121, projection='3d')
     ax2 = fig.add_subplot(122, projection='3d')
     coils_initial.plot(ax=ax1, show=False)
-    field.plot(ax=ax1, show=False)
+    field.plot(ax=ax1, show=False, alpha=0.2)
     tracing_initial.plot(ax=ax1, show=False)
     coils_optimized.plot(ax=ax2, show=False)
-    field.plot(ax=ax2, show=False)
+    field.plot(ax=ax2, show=False, alpha=0.2)
     tracing_optimized.plot(ax=ax2, show=False)
     plt.tight_layout()
     plt.show()

essos/coils.py

@@ -314,8 +314,8 @@ def __init__(self, curves: Curves, currents: jnp.ndarray):
         assert jnp.size(currents) == jnp.size(curves.dofs, 0)
         super().__init__(curves.dofs, curves.n_segments, curves.nfp, curves.stellsym)
         self._currents_scale = jnp.mean(jnp.abs(currents))
-        self._dofs_currents = currents/self.currents_scale
-        self._currents = apply_symmetries_to_currents(self._dofs_currents*self.currents_scale, self.nfp, self.stellsym)
+        self._dofs_currents = currents/self._currents_scale
+        self._currents = apply_symmetries_to_currents(self._dofs_currents*self._currents_scale, self.nfp, self.stellsym)
 
     def __str__(self):
         return f"nfp stellsym order\n{self.nfp} {self.stellsym} {self.order}\n"\

essos/dynamics.py

@@ -2,15 +2,18 @@
 jax.config.update("jax_enable_x64", True)
 import jax.numpy as jnp
 import matplotlib.pyplot as plt
-from jax.experimental.shard_map import shard_map
-from jax.sharding import Mesh, NamedSharding, PartitionSpec
-from jax import jit, vmap, tree_util, random, lax
+from jax.sharding import Mesh, PartitionSpec, NamedSharding
+from jax import jit, vmap, tree_util, random, lax, device_put
 from functools import partial
 from diffrax import diffeqsolve, ODETerm, SaveAt, Tsit5, PIDController, Event
 from essos.coils import Coils
 from essos.fields import BiotSavart, Vmec
 from essos.constants import ALPHA_PARTICLE_MASS, ALPHA_PARTICLE_CHARGE, FUSION_ALPHA_PARTICLE_ENERGY
 from essos.plot import fix_matplotlib_3d
+from essos.util import roots
+
+mesh = Mesh(jax.devices(), ("dev",))
+sharding = NamedSharding(mesh, PartitionSpec("dev", None))
 
 def gc_to_fullorbit(field, initial_xyz, initial_vparallel, total_speed, mass, charge, phase_angle_full_orbit=0):
     """
@@ -211,6 +214,7 @@ def compute_energy_fo(trajectory):
     def trace(self):
         @jit
         def compute_trajectory(initial_condition) -> jnp.ndarray:
+            # initial_condition = initial_condition[0]
             if self.model == 'FullOrbit_Boris':
                 dt=self.maxtime / self.timesteps
                 def update_state(state, _):
@@ -233,6 +237,8 @@ def update_state(state, _):
                 _, trajectory = lax.scan(update_state, initial_condition, jnp.arange(len(self.times)-1))
                 trajectory = jnp.vstack([initial_condition, trajectory])
             else:
+                import warnings
+                warnings.simplefilter("ignore", category=FutureWarning) # see https://github.com/patrick-kidger/diffrax/issues/445 for explanation
                 trajectory = diffeqsolve(
                     self.ODE_term,
                     t0=0.0,
@@ -245,15 +251,26 @@ def update_state(state, _):
                     throw=False,
                     # adjoint=DirectAdjoint(),
                     stepsize_controller = PIDController(pcoeff=0.4, icoeff=0.3, dcoeff=0, rtol=self.tol_step_size, atol=self.tol_step_size),
-                    # max_steps=1000000,
+                    max_steps=10000000000,
                     event = Event(self.condition)
                 ).ys
             return trajectory
-        # return jnp.array(vmap(compute_trajectory)(self.initial_conditions))
-        mesh = Mesh(devices=jax.devices(), axis_names=('workers',))
-        in_spec = PartitionSpec('workers', None)
-        return shard_map(vmap(compute_trajectory), mesh, 
-                        in_specs=(in_spec,), out_specs=in_spec, check_rep=False)(self.initial_conditions)
+        
+        # if len(jax.devices())!=len(self.initial_conditions):
+        #     return vmap(compute_trajectory)(self.initial_conditions[:,None,:])
+        # else:
+        #     # num_devices = len(jax.devices())
+        #     shape = self.initial_conditions.shape
+        #     # distributed_initial_conditions = self.initial_conditions.reshape(num_devices, -1, *shape[1:])
+        #     mesh = Mesh(devices=jax.devices(), axis_names=('workers'))
+        #     in_spec = PartitionSpec('workers')  # Distribute along the workers axis
+        #     out_spec = PartitionSpec('workers')  # Gather results along the same axis
+        #     return shard_map(compute_trajectory, mesh, in_specs=in_spec, out_specs=out_spec, check_rep=False)(
+        #         self.initial_conditions).reshape((shape[0], self.timesteps, shape[1]))
+        
+        return jit(vmap(compute_trajectory), in_shardings=sharding, out_shardings=sharding)(
+            device_put(self.initial_conditions, sharding))
+        
         # trajectories = []
         # for initial_condition in self.initial_conditions:
         #     trajectory = compute_trajectory(initial_condition)
@@ -289,12 +306,13 @@ def to_vtk(self, filename):
         data = np.array(jnp.concatenate([i*jnp.ones((self.trajectories[i].shape[0], )) for i in range(len(self.trajectories))]))
         polyLinesToVTK(filename, x, y, z, pointsPerLine=ppl, pointData={'idx': data})
 
-    def plot(self, ax=None, show=True, axis_equal=True, **kwargs):
+    def plot(self, ax=None, show=True, axis_equal=True, n_trajectories_plot=5, **kwargs):
         if ax is None or ax.name != "3d":
             fig = plt.figure()
             ax = fig.add_subplot(projection='3d')
         trajectories_xyz = jnp.array(self.trajectories_xyz)
-        for i in range(trajectories_xyz.shape[0]):
+        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.grid(False)
         if axis_equal:
@@ -315,101 +333,97 @@ def loss_fraction(self, r_max=0.99):
         total_particles_lost = loss_fractions[-1] * len(self.trajectories)
         return loss_fractions, total_particles_lost, lost_times
 
-    # def get_phi(x, y, phi_last):
-    #     """Compute the toroidal angle phi, ensuring continuity."""
-    #     phi = jnp.arctan2(y, x)
-    #     dphi = phi - phi_last
-    #     return phi - jnp.round(dphi / (2 * jnp.pi)) * (2 * jnp.pi)  # Ensure continuity
-
-    # @partial(jit, static_argnums=(0, 2))
-    # def find_poincare_hits(self, traj, phis_poincare):
-    #     """Find points where field lines cross specified phi values."""
-    #     x, y, z = traj[:, 0], traj[:, 1], traj[:, 2]
-    #     phi_values = jnp.unwrap(jnp.arctan2(y, x))  # Ensure continuity
-    #     t_steps = jnp.arange(len(x))
-
-    #     hits = []
-        
-    #     for phi_target in phis_poincare:
-    #         phi_shifted = phi_values - phi_target  # Shifted phi for comparison
-    #         sign_change = (phi_shifted[:-1] * phi_shifted[1:]) < 0  # Detect crossing
-
-    #         if jnp.any(sign_change):
-    #             crossing_indices = jnp.where(sign_change)[0]  # Get indices of crossings
-    #             for idx in crossing_indices:
-    #                 # Linear interpolation to estimate exact crossing
-    #                 w = (phi_target - phi_values[idx]) / (phi_values[idx + 1] - phi_values[idx])
-    #                 t_cross = t_steps[idx] + w * (t_steps[idx + 1] - t_steps[idx])
-    #                 x_cross = x[idx] + w * (x[idx + 1] - x[idx])
-    #                 y_cross = y[idx] + w * (y[idx + 1] - y[idx])
-    #                 z_cross = z[idx] + w * (z[idx + 1] - z[idx])
+    def poincare_plot(self, shifts = [jnp.pi/2], orientation = 'toroidal', length = 1, ax=None, show=True, color=None, **kwargs):
+        """
+        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].
+            orientation (str, optional): 
+                'toroidal' - find time values when toroidal angle = shift [0, 2pi].
+                'z' - find time values where z coordinate = shift. Default is 'toroidal'.
+            length (float, optional): A way to shorten data. 1 - plot full length, 0.1 - plot 1/10 of data length. Default is 1.
+            ax (matplotlib.axes._subplots.AxesSubplot, optional): Matplotlib axis to plot on. Default is None.
+            show (bool, optional): Whether to display the plot. Default is True.
+            color: Can be time, None or a color to plot Poincaré points
+            **kwargs: Additional keyword arguments for plotting.
+        Notes:
+            - If the data seem ill-behaved, there may not be enough steps in the trace for a good interpolation.
+            - This will break if there are any NaNs.
+            - Issues with toroidal interpolation: jnp.arctan2(Y, X) % (2 * jnp.pi) causes distortion in interpolation near phi = 0.
+            - Maybe determine a lower limit on resolution needed per toroidal turn for "good" results.
+        To-Do:
+            - Format colorbars.
+        """
+        kwargs.setdefault('s', 0.5)
+        if ax is None:
+            fig = plt.figure()
+            ax = fig.add_subplot()
+        shifts = jnp.array(shifts)
+        plotting_data = []
+        # from essos.util import roots_scipy
+        for shift in shifts:
+            @jit
+            def compute_trajectory_toroidal(trace):
+                X,Y,Z = trace[:,:3].T
+                R = jnp.sqrt(X**2 + Y**2)
+                phi = jnp.arctan2(Y,X)
+                phi = jnp.where(shift==0, phi, jnp.abs(phi))
+                T_slice = roots(self.times, phi, shift = shift)
+                T_slice = jnp.where(shift==0, jnp.concatenate((T_slice[1::2],T_slice[1::2])), T_slice)
+                # T_slice = roots_scipy(self.times, phi, shift = shift)
+                R_slice = jnp.interp(T_slice, self.times, R)
+                Z_slice = jnp.interp(T_slice, self.times, Z)
+                return R_slice, Z_slice, T_slice
+            @jit
+            def compute_trajectory_z(trace):
+                X,Y,Z = trace[:,:3].T
+                T_slice = roots(self.times, Z, shift = shift)
+                # T_slice = roots_scipy(self.times, Z, shift = shift)
+                X_slice = jnp.interp(T_slice, self.times, X)
+                Y_slice = jnp.interp(T_slice, self.times, Y)
+                return X_slice, Y_slice, T_slice
+            if orientation == 'toroidal':
+                # X_slice, Y_slice, T_slice = vmap(compute_trajectory_toroidal)(self.trajectories)
+                X_slice, Y_slice, T_slice = jit(vmap(compute_trajectory_toroidal), in_shardings=sharding, out_shardings=sharding)(
+                    device_put(self.trajectories, sharding))
+            elif orientation == 'z':
+                # X_slice, Y_slice, T_slice = vmap(compute_trajectory_z)(self.trajectories)
+                X_slice, Y_slice, T_slice = jit(vmap(compute_trajectory_z), in_shardings=sharding, out_shardings=sharding)(
+                    device_put(self.trajectories, sharding))
+            @partial(jax.vmap, in_axes=(0, 0, 0))
+            def process_trajectory(X_i, Y_i, T_i):
+                mask = (T_i[1:] != T_i[:-1])
+                valid_idx = jnp.nonzero(mask, size=T_i.size - 1)[0] + 1
+                return X_i[valid_idx], Y_i[valid_idx], T_i[valid_idx]
+            X_s, Y_s, T_s = process_trajectory(X_slice, Y_slice, T_slice)
+            length_ = (vmap(len)(X_s) * length).astype(int)
+            colors = plt.cm.ocean(jnp.linspace(0, 0.8, len(X_s)))
+            for i in range(len(X_s)):
+                X_plot, Y_plot = X_s[i][:length_[i]], Y_s[i][:length_[i]]
+                T_plot = T_s[i][:length_[i]]
+                plotting_data.append((X_plot, Y_plot, T_plot))
+                if color == 'time':
+                    hits = ax.scatter(X_plot, Y_plot, c=T_s[i][:length_[i]], **kwargs)
+                else:
+                    if color is None: c=[colors[i]]
+                    else: c=color
+                    hits = ax.scatter(X_plot, Y_plot, c=c, **kwargs)
 
-    #                 hits.append([t_cross, x_cross, y_cross, z_cross])
-
-    #     return jnp.array(hits)
-
-    # @partial(jit, static_argnums=(0))
-    # def poincare(self):
-    #     """Compute Poincaré section hits for multiple trajectories."""
-    #     trajectories = self.trajectories  # Pass trajectories directly into the function
-    #     phis_poincare = self.phis_poincare  # Similarly, use the direct attribute
-
-    #     # Use vmap to vectorize the calls for each trajectory
-    #     return vmap(self.find_poincare_hits, in_axes=(0, None))(trajectories, tuple(phis_poincare))
-
-    # def poincare_plot(self, phis=None, filename=None, res_phi_hits=None, mark_lost=False, aspect='equal', dpi=300, xlims=None, 
-    #                 ylims=None, s=2, marker='o', show=True):
-    #     import matplotlib.pyplot as plt
-        
-    #     self.phis_poincare = phis
-    #     if res_phi_hits is None:
-    #         res_phi_hits = self.poincare()
-    #     self.res_phi_hits = res_phi_hits
-            
-    #     res_phi_hits = jnp.array(res_phi_hits)  # Ensure it's a JAX array
-        
-    #     # Determine number of rows/columns
-    #     nrowcol = int(jnp.ceil(jnp.sqrt(len(phis))))
-        
-    #     # Create subplots
-    #     fig, axs = plt.subplots(nrowcol, nrowcol, figsize=(8, 5))
-    #     axs = axs.ravel()  # Flatten for easier indexing
+        if orientation == 'toroidal':
+            plt.xlabel('R',fontsize = 18)
+            plt.ylabel('Z',fontsize = 18)
+            # plt.title(r'$\phi$ = {:.2f} $\pi$'.format(shift/jnp.pi),fontsize = 20)
+        elif orientation == 'z':
+            plt.xlabel('X',fontsize = 18)
+            plt.xlabel('Y',fontsize = 18)
+            # plt.title('Z = {:.2f}'.format(shift),fontsize = 20)
+        plt.axis('equal')
+        plt.grid()
+        plt.tight_layout()
+        if show:
+            plt.show()
 
-    #     # Loop over phi values and create plots
-    #     for i, phi in enumerate(phis):
-    #         ax = axs[i]
-    #         ax.set_aspect(aspect)
-    #         ax.set_title(f"$\\phi = {phi/jnp.pi:.2f}\\pi$", loc='left', y=0.0)
-    #         ax.set_xlabel("$r$")
-    #         ax.set_ylabel("$z$")
-            
-    #         if xlims:
-    #             ax.set_xlim(xlims)
-    #         if ylims:
-    #             ax.set_ylim(ylims)
-            
-    #         # Extract points corresponding to this phi
-    #         mask = res_phi_hits[:, 1] == i
-    #         data_this_phi = res_phi_hits[mask]
-            
-    #         if data_this_phi.shape[0] > 0:
-    #             r = jnp.sqrt(data_this_phi[:, 2]**2 + data_this_phi[:, 3]**2)
-    #             z = data_this_phi[:, 4]
-                
-    #             color = 'g'  # Default color
-    #             if mark_lost:
-    #                 lost = data_this_phi[-1, 1] < 0
-    #                 color = 'r' if lost else 'g'
-                    
-    #             ax.scatter(r, z, marker=marker, s=s, linewidths=0, c=color)
-
-    #         ax.grid(True, linewidth=0.5)
-
-    #     # Adjust layout and save
-    #     plt.tight_layout()
-    #     if filename is not None: plt.savefig(filename, dpi=dpi)
-    #     if show: plt.show()
-    #     plt.close()
+        return plotting_data
 
 tree_util.register_pytree_node(Tracing,
                                Tracing._tree_flatten,