Merge pull request #385 from hiddenSymmetries/rz_stopping_criteria

RZ stopping criteria

src/simsopt/field/tracing.py

@@ -16,7 +16,9 @@
 logger = logging.getLogger(__name__)
 
 __all__ = ['SurfaceClassifier', 'LevelsetStoppingCriterion',
-           'MinToroidalFluxStoppingCriterion', 'MaxToroidalFluxStoppingCriterion',
+           'MinToroidalFluxStoppingCriterion','MaxToroidalFluxStoppingCriterion',
+           'MinRStoppingCriterion','MinZStoppingCriterion',
+           'MaxRStoppingCriterion','MaxZStoppingCriterion',
            'IterationStoppingCriterion', 'ToroidalTransitStoppingCriterion',
            'compute_fieldlines', 'compute_resonances',
            'compute_poloidal_transits', 'compute_toroidal_transits',
@@ -738,7 +740,6 @@ def __init__(self, classifier):
         else:
             sopp.LevelsetStoppingCriterion.__init__(self, classifier)
 
-
 class MinToroidalFluxStoppingCriterion(sopp.MinToroidalFluxStoppingCriterion):
     """
     Stop the iteration once a particle falls below a critical value of
@@ -774,7 +775,6 @@ class MaxToroidalFluxStoppingCriterion(sopp.MaxToroidalFluxStoppingCriterion):
     """
     pass
 
-
 class ToroidalTransitStoppingCriterion(sopp.ToroidalTransitStoppingCriterion):
     """
     Stop the iteration once the maximum number of toroidal transits is reached.
@@ -797,6 +797,65 @@ class IterationStoppingCriterion(sopp.IterationStoppingCriterion):
     """
     pass
 
+class MinRStoppingCriterion(sopp.MinRStoppingCriterion):
+    """
+    Stop the iteration once a particle falls below a critical value of
+    ``R``, the radial cylindrical coordinate. 
+
+    Usage:
+
+    .. code-block::
+
+        stopping_criteria=[MinRStopingCriterion(crit_r)]
+
+    where ``crit_r`` is the value of the critical coordinate.
+    """
+    pass
+
+class MinZStoppingCriterion(sopp.MinZStoppingCriterion):
+    """
+    Stop the iteration once a particle falls below a critical value of
+    ``Z``, the cylindrical vertical coordinate. 
+
+    Usage:
+
+    .. code-block::
+
+        stopping_criteria=[MinZStopingCriterion(crit_z)]
+
+    where ``crit_z`` is the value of the critical coordinate.
+    """
+    pass
+
+class MaxRStoppingCriterion(sopp.MaxRStoppingCriterion):
+    """
+    Stop the iteration once a particle goes above a critical value of
+    ``R``, the radial cylindrical coordinate. 
+
+    Usage:
+
+    .. code-block::
+
+        stopping_criteria=[MaxRStopingCriterion(crit_r)]
+
+    where ``crit_r`` is the value of the critical coordinate.
+    """
+    pass
+
+class MaxZStoppingCriterion(sopp.MaxZStoppingCriterion):
+    """
+    Stop the iteration once a particle gove above a critical value of
+    ``Z``, the cylindrical vertical coordinate. 
+
+    Usage:
+
+    .. code-block::
+
+        stopping_criteria=[MaxZStopingCriterion(crit_z)]
+
+    where ``crit_z`` is the value of the critical coordinate.
+    """
+    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,6 +17,14 @@ 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_<MinRStoppingCriterion, shared_ptr<MinRStoppingCriterion>, StoppingCriterion>(m, "MinRStoppingCriterion")
+        .def(py::init<double>());
+    py::class_<MinZStoppingCriterion, shared_ptr<MinZStoppingCriterion>, StoppingCriterion>(m, "MinZStoppingCriterion")
+        .def(py::init<double>());
+    py::class_<MaxRStoppingCriterion, shared_ptr<MaxRStoppingCriterion>, StoppingCriterion>(m, "MaxRStoppingCriterion")
+        .def(py::init<double>());
+    py::class_<MaxZStoppingCriterion, shared_ptr<MaxZStoppingCriterion>, StoppingCriterion>(m, "MaxZStoppingCriterion")
+        .def(py::init<double>());
     py::class_<MaxToroidalFluxStoppingCriterion, shared_ptr<MaxToroidalFluxStoppingCriterion>, StoppingCriterion>(m, "MaxToroidalFluxStoppingCriterion")
         .def(py::init<double>());
     py::class_<MinToroidalFluxStoppingCriterion, shared_ptr<MinToroidalFluxStoppingCriterion>, StoppingCriterion>(m, "MinToroidalFluxStoppingCriterion")

src/simsoptpp/tracing.h

@@ -64,6 +64,46 @@ class MinToroidalFluxStoppingCriterion : public StoppingCriterion{
         };
 };
 
+class MinZStoppingCriterion : public StoppingCriterion{
+    private:
+        double crit_z;
+    public:
+        MinZStoppingCriterion(double crit_z) : crit_z(crit_z) {};
+        bool operator()(int iter, double t, double x, double y, double z) override {
+            return z<=crit_z;
+        };
+};
+
+class MaxZStoppingCriterion : public StoppingCriterion{
+    private:
+        double crit_z;
+    public:
+        MaxZStoppingCriterion(double crit_z) : crit_z(crit_z) {};
+        bool operator()(int iter, double t, double x, double y, double z) override {
+            return z>=crit_z;
+        };
+};
+
+class MinRStoppingCriterion : public StoppingCriterion{
+    private:
+        double crit_r;
+    public:
+        MinRStoppingCriterion(double crit_r) : crit_r(crit_r) {};
+        bool operator()(int iter, double t, double x, double y, double z) override {
+            return std::sqrt(x*x+y*y)<=crit_r;            
+        };
+};
+
+class MaxRStoppingCriterion : public StoppingCriterion{
+    private:
+        double crit_r;
+    public:
+        MaxRStoppingCriterion(double crit_r) : crit_r(crit_r) {};
+        bool operator()(int iter, double t, double x, double y, double z) override {
+            return std::sqrt(x*x+y*y)>=crit_r;            
+        };
+};
+
 class IterationStoppingCriterion : public StoppingCriterion{
     private:
         int max_iter;

tests/field/test_fieldline.py

@@ -3,7 +3,8 @@
 import numpy as np
 
 from simsopt.field.magneticfieldclasses import ToroidalField, PoloidalField, InterpolatedField, UniformInterpolationRule
-from simsopt.field.tracing import compute_fieldlines, particles_to_vtk, plot_poincare_data
+from simsopt.field.tracing import compute_fieldlines, particles_to_vtk, plot_poincare_data, \
+    MinRStoppingCriterion, MinZStoppingCriterion, MaxRStoppingCriterion, MaxZStoppingCriterion
 from simsopt.field.biotsavart import BiotSavart
 from simsopt.configs.zoo import get_ncsx_data
 from simsopt.field.coil import coils_via_symmetries, Coil, Current
@@ -136,7 +137,36 @@ def test_poincare_caryhanson(self):
         Z0 = [ma.gamma()[0, 2]]        
         phis = np.arctan2(ma.gamma()[:, 1], ma.gamma()[:, 0])
         res_tys, res_phi_hits = compute_fieldlines(
-            bs, R0, Z0, tmax=2, phis=phis, stopping_criteria=[])
-
+            bs, R0, Z0, tmax=2, phis=phis)
         for i in range(len(res_phi_hits[0])):
             assert np.linalg.norm(ma.gamma()[i+1, :] - res_phi_hits[0][i, 2:5]) < 2e-3
+
+        # Text StoppingCriterion in R and Z 
+        # For each case, check that stopping criterion was met. 
+        # Check that R/Z is less than/greater than the maximum/minimum value. 
+        Rmax = 1
+        res_tys, res_phi_hits = compute_fieldlines(
+            bs, [Rmax-0.02], [1], tmax=2000, stopping_criteria=[MaxRStoppingCriterion(Rmax)])
+        assert res_phi_hits[0][0,1] == -1
+        assert np.all(np.sqrt(res_tys[0][:, 1]**2 + res_tys[0][:, 2]**2) < Rmax)
+
+        Rmin = 0.3
+        res_tys, res_phi_hits = compute_fieldlines(
+            bs, [Rmin+0.02], [0.3], tmax=500, stopping_criteria=[MinRStoppingCriterion(Rmin)])
+        assert res_phi_hits[0][0,1] == -1
+        assert np.all(np.sqrt(res_tys[0][:, 1]**2 + res_tys[0][:, 2]**2) > Rmin) 
+
+        Zmin = -0.1
+        res_tys, res_phi_hits = compute_fieldlines(
+            bs, [0.97], [Zmin+0.02], tmax=2000, stopping_criteria=[MinZStoppingCriterion(Zmin)]
+            )
+        assert res_phi_hits[0][0,1] == -1
+        assert np.all(res_tys[0][:, 3] > Zmin)
+
+        Zmax = 0.5
+        res_tys, res_phi_hits = compute_fieldlines(
+            bs, [0.5], [Zmax-0.02], tmax=2000, stopping_criteria=[MaxZStoppingCriterion(Zmax)]
+            )
+        assert res_phi_hits[0][0,1] == -1
+        assert np.all(res_tys[0][:, 3] < Zmax)
+