RZ stopping criteria

examples/2_Intermediate/tracing_boozer.py

@@ -6,7 +6,7 @@
 import matplotlib.pyplot as plt
 
 from simsopt.field import (BoozerRadialInterpolant, InterpolatedBoozerField, trace_particles_boozer,
-                           MinToroidalFluxStoppingCriterion, MaxToroidalFluxStoppingCriterion,
+                           ToroidalFluxStoppingCriterion,
                            ToroidalTransitStoppingCriterion, compute_resonances)
 from simsopt.mhd import Vmec
 from simsopt.util import in_github_actions
@@ -75,7 +75,7 @@
 
 gc_tys, gc_zeta_hits = trace_particles_boozer(
     field, stz_inits, vpar_inits, tmax=1e-2, mass=mass, charge=ELEMENTARY_CHARGE,
-    Ekin=Ekin, tol=1e-8, mode='gc_vac', stopping_criteria=[MaxToroidalFluxStoppingCriterion(0.99), MinToroidalFluxStoppingCriterion(0.01), ToroidalTransitStoppingCriterion(100, True)],
+    Ekin=Ekin, tol=1e-8, mode='gc_vac', stopping_criteria=[ToroidalFluxStoppingCriterion(0.99,False), ToroidalFluxStoppingCriterion(0.01,True), ToroidalTransitStoppingCriterion(100, True)],
     forget_exact_path=False)
 
 Nparticles = len(gc_tys)
@@ -112,7 +112,7 @@
 
 gc_tys, gc_zeta_hits = trace_particles_boozer(
     field, stz_inits, vpar_inits, tmax=1e-2, mass=mass, charge=ELEMENTARY_CHARGE,
-    Ekin=Ekin, zetas=[0], tol=1e-8, stopping_criteria=[MinToroidalFluxStoppingCriterion(0.01), MaxToroidalFluxStoppingCriterion(0.99), ToroidalTransitStoppingCriterion(30, True)],
+    Ekin=Ekin, zetas=[0], tol=1e-8, stopping_criteria=[ToroidalFluxStoppingCriterion(0.01,True), ToroidalFluxStoppingCriterion(0.99,False), ToroidalTransitStoppingCriterion(30, True)],
     forget_exact_path=False)
 
 resonances = compute_resonances(gc_tys, gc_zeta_hits, ma=None, delta=0.01)
@@ -134,7 +134,7 @@
 
 gc_tys, gc_zeta_hits = trace_particles_boozer(
     field, stz_res, vpar_res, tmax=1e-2, mass=mass, charge=ELEMENTARY_CHARGE,
-    Ekin=Ekin, zetas=[0], tol=1e-8, stopping_criteria=[MinToroidalFluxStoppingCriterion(0.01), MaxToroidalFluxStoppingCriterion(0.99), ToroidalTransitStoppingCriterion(30, True)],
+    Ekin=Ekin, zetas=[0], tol=1e-8, stopping_criteria=[ToroidalFluxStoppingCriterion(0.01,True), ToroidalFluxStoppingCriterion(0.99,False), ToroidalTransitStoppingCriterion(30, True)],
     forget_exact_path=False)
 
 if not in_github_actions:

src/simsopt/field/tracing.py

@@ -16,7 +16,7 @@
 logger = logging.getLogger(__name__)
 
 __all__ = ['SurfaceClassifier', 'LevelsetStoppingCriterion',
-           'MinToroidalFluxStoppingCriterion', 'MaxToroidalFluxStoppingCriterion',
+           'ToroidalFluxStoppingCriterion','RStoppingCriterion','ZStoppingCriterion',
            'IterationStoppingCriterion', 'ToroidalTransitStoppingCriterion',
            'compute_fieldlines', 'compute_resonances',
            'compute_poloidal_transits', 'compute_toroidal_transits',
@@ -738,66 +738,87 @@ def __init__(self, classifier):
         else:
             sopp.LevelsetStoppingCriterion.__init__(self, classifier)
 
-
-class MinToroidalFluxStoppingCriterion(sopp.MinToroidalFluxStoppingCriterion):
+class ToroidalFluxStoppingCriterion(sopp.ToroidalFluxStoppingCriterion):
     """
-    Stop the iteration once a particle falls below a critical value of
+    Stop the iteration once a particle falls above or below a critical value of
     ``s``, the normalized toroidal flux. This :class:`StoppingCriterion` is
-    important to use when tracing particles in flux coordinates, as the poloidal
-    angle becomes ill-defined at the magnetic axis. This should only be used
-    when tracing trajectories in a flux coordinate system (i.e., :class:`trace_particles_boozer`).
+    important to use when tracing particles in flux coordinates. For example,
+    the poloidal angle becomes ill-defined at the magnetic axis. This should 
+    only be used when tracing trajectories in a flux coordinate system (i.e., 
+    :class:`trace_particles_boozer`).
 
     Usage:
 
     .. code-block::
 
-        stopping_criteria=[MinToroidalFluxStopingCriterion(s)]
+        stopping_criteria=[ToroidalFluxStopingCriterion(crit_s,min_bool)]
 
-    where ``s`` is the value of the minimum normalized toroidal flux.
+    where ``crit_s`` is the value of the critical normalized toroidal flux and 
+    ``min_bool'' is a boolean indicating whether to stop when  
+    ``s'' is less than the critical value (``true'') or 
+    greater than the critical value (``false'').
     """
     pass
 
 
-class MaxToroidalFluxStoppingCriterion(sopp.MaxToroidalFluxStoppingCriterion):
+class ToroidalTransitStoppingCriterion(sopp.ToroidalTransitStoppingCriterion):
     """
-    Stop the iteration once a particle falls above a critical value of
-    ``s``, the normalized toroidal flux. This should only be used when tracing
-    trajectories in a flux coordinate system (i.e., :class:`trace_particles_boozer`).
+    Stop the iteration once the maximum number of toroidal transits is reached.
 
     Usage:
 
     .. code-block::
 
-        stopping_criteria=[MaxToroidalFluxStopingCriterion(s)]
+        stopping_criteria=[ToroidalTransitStoppingCriterion(ntransits,flux)]
 
-    where ``s`` is the value of the maximum normalized toroidal flux.
+    where ``ntransits`` is the maximum number of toroidal transits and ``flux``
+    is a boolean indicating whether tracing is being performed in a flux coordinate system.
     """
     pass
 
 
-class ToroidalTransitStoppingCriterion(sopp.ToroidalTransitStoppingCriterion):
+class IterationStoppingCriterion(sopp.IterationStoppingCriterion):
     """
-    Stop the iteration once the maximum number of toroidal transits is reached.
+    Stop the iteration once the maximum number of iterations is reached.
+    """
+    pass
+
+class RStoppingCriterion(sopp.RStoppingCriterion):
+    """
+    Stop the iteration once a particle falls above or below a critical value of
+    ``R``, the radial cylindrical coordinate. 
 
     Usage:
 
     .. code-block::
 
-        stopping_criteria=[ToroidalTransitStoppingCriterion(ntransits,flux)]
+        stopping_criteria=[RStopingCriterion(crit_r,min_bool)]
 
-    where ``ntransits`` is the maximum number of toroidal transits and ``flux``
-    is a boolean indicating whether tracing is being performed in a flux coordinate system.
+    where ``crit_r`` is the value of the critical coordinate and 
+    ``min_bool'' is a boolean indicating whether to stop when  
+    ``R'' is less than the critical value (``true'') or 
+    greater than the critical value (``false'').
     """
     pass
 
-
-class IterationStoppingCriterion(sopp.IterationStoppingCriterion):
+class ZStoppingCriterion(sopp.ZStoppingCriterion):
     """
-    Stop the iteration once the maximum number of iterations is reached.
+    Stop the iteration once a particle falls above or below a critical value of
+    ``Z``, the cylindrical vertical coordinate. 
+
+    Usage:
+
+    .. code-block::
+
+        stopping_criteria=[ZStopingCriterion(crit_z,min_bool)]
+
+    where ``crit_z`` is the value of the critical coordinate and 
+    ``min_bool'' is a boolean indicating whether to stop when  
+    ``Z'' is less than the critical value (``true'') or 
+    greater than the critical value (``false'').
     """
     pass
 
-
 def plot_poincare_data(fieldlines_phi_hits, phis, filename, mark_lost=False, aspect='equal', dpi=300, xlims=None, 
                        ylims=None, surf=None, s=2, marker='o'):
     """

src/simsoptpp/python_tracing.cpp

@@ -17,10 +17,12 @@ void init_tracing(py::module_ &m){
     py::class_<StoppingCriterion, shared_ptr<StoppingCriterion>>(m, "StoppingCriterion");
     py::class_<IterationStoppingCriterion, shared_ptr<IterationStoppingCriterion>, StoppingCriterion>(m, "IterationStoppingCriterion")
         .def(py::init<int>());
-    py::class_<MaxToroidalFluxStoppingCriterion, shared_ptr<MaxToroidalFluxStoppingCriterion>, StoppingCriterion>(m, "MaxToroidalFluxStoppingCriterion")
-        .def(py::init<double>());
-    py::class_<MinToroidalFluxStoppingCriterion, shared_ptr<MinToroidalFluxStoppingCriterion>, StoppingCriterion>(m, "MinToroidalFluxStoppingCriterion")
-        .def(py::init<double>());
+    py::class_<RStoppingCriterion, shared_ptr<RStoppingCriterion>, StoppingCriterion>(m, "RStoppingCriterion")
+        .def(py::init<double,bool>());
+    py::class_<ZStoppingCriterion, shared_ptr<ZStoppingCriterion>, StoppingCriterion>(m, "ZStoppingCriterion")
+        .def(py::init<double,bool>());
+    py::class_<ToroidalFluxStoppingCriterion, shared_ptr<ToroidalFluxStoppingCriterion>, StoppingCriterion>(m, "ToroidalFluxStoppingCriterion")
+        .def(py::init<double,bool>());
     py::class_<ToroidalTransitStoppingCriterion, shared_ptr<ToroidalTransitStoppingCriterion>, StoppingCriterion>(m, "ToroidalTransitStoppingCriterion")
         .def(py::init<int,bool>());
     py::class_<LevelsetStoppingCriterion<PyTensor>, shared_ptr<LevelsetStoppingCriterion<PyTensor>>, StoppingCriterion>(m, "LevelsetStoppingCriterion")

src/simsoptpp/tracing.h

@@ -44,23 +44,51 @@ class ToroidalTransitStoppingCriterion : public StoppingCriterion {
         };
 };
 
-class MaxToroidalFluxStoppingCriterion : public StoppingCriterion{
+class ToroidalFluxStoppingCriterion : public StoppingCriterion{
     private:
-        double max_s;
+        double crit_s;
+        bool min_bool;
     public:
-        MaxToroidalFluxStoppingCriterion(double max_s) : max_s(max_s) {};
+        ToroidalFluxStoppingCriterion(double crit_s, bool min_bool) : crit_s(crit_s), min_bool(min_bool) {};
         bool operator()(int iter, double t, double s, double theta, double zeta) override {
-            return s>=max_s;
+            if (min_bool) {
+                return s<=crit_s;
+            } else {
+                return s>=crit_s;
+            }
+
         };
 };
 
-class MinToroidalFluxStoppingCriterion : public StoppingCriterion{
+class ZStoppingCriterion : public StoppingCriterion{
     private:
-        double min_s;
+        double crit_z;
+        bool min_bool;
     public:
-        MinToroidalFluxStoppingCriterion(double min_s) : min_s(min_s) {};
-        bool operator()(int iter, double t, double s, double theta, double zeta) override {
-            return s<=min_s;
+        ZStoppingCriterion(double crit_z, bool min_bool) : crit_z(crit_z), min_bool(min_bool) {};
+        bool operator()(int iter, double t, double x, double y, double z) override {
+            if (min_bool) {
+                return z<=crit_z;
+            } else {
+                return z>=crit_z;
+            }
+
+        };
+};
+
+class RStoppingCriterion : public StoppingCriterion{
+    private:
+        double crit_r;
+        bool min_bool;
+    public:
+        RStoppingCriterion(double crit_r, bool min_bool) : crit_r(crit_r), min_bool(min_bool) {};
+        bool operator()(int iter, double t, double x, double y, double z) override {
+            if (min_bool) {
+                return std::sqrt(x*x+y*y)<=crit_r;
+            } else {
+                return std::sqrt(x*x+y*y)>=crit_r;
+            }
+
         };
 };
 

tests/field/test_particle.py

@@ -12,7 +12,7 @@
 from simsopt.field.tracing import trace_particles_starting_on_curve, SurfaceClassifier, \
     particles_to_vtk, LevelsetStoppingCriterion, compute_gc_radius, gc_to_fullorbit_initial_guesses, \
     IterationStoppingCriterion, trace_particles_starting_on_surface, trace_particles_boozer, \
-    MinToroidalFluxStoppingCriterion, MaxToroidalFluxStoppingCriterion, ToroidalTransitStoppingCriterion, \
+    ToroidalFluxStoppingCriterion, ToroidalTransitStoppingCriterion, \
     compute_poloidal_transits, compute_toroidal_transits, trace_particles, compute_resonances
 from simsopt.geo.surfacerzfourier import SurfaceRZFourier
 from simsopt.field.boozermagneticfield import BoozerAnalytic
@@ -451,7 +451,7 @@ def test_energy_momentum_conservation_boozer(self):
 
         gc_tys, gc_phi_hits = trace_particles_boozer(bsh, stz_inits, vpar_inits,
                                                      tmax=tmax, mass=m, charge=q, Ekin=Ekin, zetas=[], mode='gc_vac',
-                                                     stopping_criteria=[MinToroidalFluxStoppingCriterion(.01), MaxToroidalFluxStoppingCriterion(0.99), ToroidalTransitStoppingCriterion(100, True)],
+                                                     stopping_criteria=[ToroidalFluxStoppingCriterion(.01,True), ToroidalFluxStoppingCriterion(0.99,False), ToroidalTransitStoppingCriterion(100, True)],
                                                      tol=1e-12)
 
         # pick 100 random points on each trace, and ensure that
@@ -515,7 +515,7 @@ def test_energy_momentum_conservation_boozer(self):
 
         gc_tys, gc_phi_hits = trace_particles_boozer(bsh, stz_inits, vpar_inits,
                                                      tmax=tmax, mass=m, charge=q, Ekin=Ekin, zetas=[], mode='gc_noK',
-                                                     stopping_criteria=[MinToroidalFluxStoppingCriterion(.01), MaxToroidalFluxStoppingCriterion(0.99), ToroidalTransitStoppingCriterion(100, True)],
+                                                     stopping_criteria=[ToroidalFluxStoppingCriterion(.01,True), ToroidalFluxStoppingCriterion(0.99,False), ToroidalTransitStoppingCriterion(100, True)],
                                                      tol=1e-12)
 
         max_energy_gc_error = np.array([])
@@ -567,7 +567,7 @@ def test_energy_momentum_conservation_boozer(self):
 
         gc_tys, gc_phi_hits = trace_particles_boozer(bsh, stz_inits, vpar_inits,
                                                      tmax=tmax, mass=m, charge=q, Ekin=Ekin, zetas=[], mode='gc',
-                                                     stopping_criteria=[MinToroidalFluxStoppingCriterion(.01), MaxToroidalFluxStoppingCriterion(0.99), ToroidalTransitStoppingCriterion(100, True)],
+                                                     stopping_criteria=[ToroidalFluxStoppingCriterion(.01,True), ToroidalFluxStoppingCriterion(0.99,False), ToroidalTransitStoppingCriterion(100, True)],
                                                      tol=1e-12)
 
         max_energy_gc_error = np.array([])
@@ -607,7 +607,7 @@ def test_energy_momentum_conservation_boozer(self):
         # Now trace with forget_exact_path = False. Check that gc_phi_hits is the same
         gc_tys, gc_phi_hits_2 = trace_particles_boozer(bsh, stz_inits, vpar_inits,
                                                        tmax=tmax, mass=m, charge=q, Ekin=Ekin, zetas=[], mode='gc_noK',
-                                                       stopping_criteria=[MinToroidalFluxStoppingCriterion(.01), MaxToroidalFluxStoppingCriterion(0.99), ToroidalTransitStoppingCriterion(100, True)],
+                                                       stopping_criteria=[ToroidalFluxStoppingCriterion(.01,True), ToroidalFluxStoppingCriterion(0.99,False), ToroidalTransitStoppingCriterion(100, True)],
                                                        tol=1e-12, forget_exact_path=True)
 
         for i in range(len(gc_phi_hits_2)):
@@ -652,7 +652,7 @@ def test_compute_poloidal_toroidal_transits(self):
 
         gc_tys, gc_phi_hits = trace_particles_boozer(bsh, stz_inits, vpar_inits,
                                                      tmax=tmax, mass=m, charge=q, Ekin=Ekin, zetas=[], mode='gc_vac',
-                                                     stopping_criteria=[MinToroidalFluxStoppingCriterion(.01), MaxToroidalFluxStoppingCriterion(0.99), ToroidalTransitStoppingCriterion(1, True)],
+                                                     stopping_criteria=[ToroidalFluxStoppingCriterion(.01,True), ToroidalFluxStoppingCriterion(0.99,False), ToroidalTransitStoppingCriterion(1, True)],
                                                      tol=1e-12)
 
         mpol = compute_poloidal_transits(gc_tys)
@@ -726,7 +726,7 @@ def test_toroidal_flux_stopping_criterion(self):
 
         gc_tys, gc_phi_hits = trace_particles_boozer(bsh, stz_inits, vpar_inits,
                                                      tmax=tmax, mass=m, charge=q, Ekin=Ekin, zetas=[], mode='gc_vac',
-                                                     stopping_criteria=[MinToroidalFluxStoppingCriterion(0.4), MaxToroidalFluxStoppingCriterion(0.6)],
+                                                     stopping_criteria=[ToroidalFluxStoppingCriterion(0.4,True), ToroidalFluxStoppingCriterion(0.6,False)],
                                                      tol=1e-12)
 
         for i in range(Nparticles):
@@ -768,7 +768,7 @@ def test_compute_resonances(self):
 
         gc_tys, gc_phi_hits = trace_particles_boozer(bsh, stz_inits, vpar_inits,
                                                      tmax=tmax, mass=m, charge=q, Ekin=Ekin, zetas=[0], mode='gc_vac',
-                                                     stopping_criteria=[MinToroidalFluxStoppingCriterion(0.01), MaxToroidalFluxStoppingCriterion(0.99), ToroidalTransitStoppingCriterion(100, True)],
+                                                     stopping_criteria=[ToroidalFluxStoppingCriterion(0.01,True), ToroidalFluxStoppingCriterion(0.99,False), ToroidalTransitStoppingCriterion(100, True)],
                                                      tol=1e-8)
 
         resonances = compute_resonances(gc_tys, gc_phi_hits, delta=0.01)