Merge pull request #236 from hiddenSymmetries/master

Update branch

conda.recipe/conda_build_config.yaml

@@ -1,12 +1,7 @@
-numpy:
-  - 1.19
-  - 1.20
-  - 1.21
 python:
   - 3.7
   - 3.8
   - 3.9
 
 pin_run_as_build:
-  numpy: x.x
   python: x.x

conda.recipe/meta.yaml

@@ -29,13 +29,13 @@ requirements:
     - pip
     - wheel
     - setuptools
-    - numpy {{ numpy }}
+    - numpy 
     - "setuptools_scm[toml]"
     - python {{ python }}
 
   run:
     - python
-    - numpy {{ numpy }}
+    - numpy {{ pin_compatible("numpy") }}
     - jax >=0.2.5
     - jaxlib >=0.1.56
     - scipy >=1.5.4

docs/source/cpp.rst

@@ -0,0 +1,65 @@
+C++ backend
+***********
+
+``SIMSOPT`` uses C++ for performance critical functions such as the Biot Savart law, many of the geometric classes, and particle tracing.
+This section is aimed at advanced developers of ``SIMSOPT`` to give an overview over the interface between C++ and Python and to help avoid common pitfalls. For most users of ``SIMSOPT`` this is not relevant.
+
+The C++ code can be found in the folder ``src/simsoptpp``.
+
+
+pybind11
+^^^^^^^^
+
+To expose the C++ code to python, we use the 
+`pybind11 <https://github.com/pybind/pybind11>`_ library.
+The interfacing happens in the ``python.cpp`` and ``python_*.cpp`` files.
+
+Trampoline classes:
+In many cases we define some parent class in C++ that has virtual functions and then we want to inherit from this class on the python side and overload these functions.
+A good example for this is the :mod:`simsoptpp.MagneticField` class. This is the base class for magnetic fields and it takes care of things like caching, or coordinate systems for evaluating magnetic fields. When you create a new magnetic field, you overload functions such as ``B_impl``, in order to compute the magnetic field at given locations. In order for this overload to work on the python side, pybind requires so called `trampoline classes <https://pybind11-jagerman.readthedocs.io/en/latest/advanced/classes.html#overriding-virtual-functions-in-python>`_. In the case of magnetic fields, this can be found in ``src/simsoptpp/pymagneticfield.h``.
+
+
+Lifetime of objects:
+memory management in hybrid Python/C++ codes can be difficult. To guarantee that objects managed by C++ are not deleted even though they are still used on the Python side, we make sure that we hold a reference to them in the Python code. A good example for this is the :mod:`simsopt.field.biotsavart.BiotSavart` class that keeps a reference to the underlying coils. See `this pull request <https://github.com/hiddenSymmetries/simsopt/pull/147>`_ for a discussion of this issue and further references.
+
+xtensor and xtensor-python
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+We use ``numpy`` array heavily on the python side; in order to use these (without copying them) in the C++ code, we employ the `xtensor <https://github.com/xtensor-stack/xtensor>`_ library and its `xtensor-python <https://github.com/xtensor-stack/xtensor-python>`_ interface to python.
+
+``xarray`` vs ``pyarray``:
+
+One pitfall to be aware of here is the following: there are different types for arrays that are managed from the python side, vs those managed on the C++ side. This means that for code that is purely C++, one should use ``xarray`` from ``xtensor``, but for code that is used from python, one uses ``pyarray`` from ``xtensor-python``. While most of ``simsoptpp`` is only ever used from the python side, it can be useful for profiling purposes (or for future usecases of the library by other C++ codes) to allow for the use of ``xarray``. For this reason, many functions and classes are templated on the data type, so that either is possible.
+
+
+OpenMP
+^^^^^^
+Some of the code is parallelized using OpenMP. OpenMP can be turned of by setting
+``export OMP_NUM_THREADS=1``
+before running a ``SIMSOPT`` script. This is recommended when debugging bugs that are assumed to be in the C++ code. We have found that creating new ``xtensor`` arrays or tensors in OpenMP threads leads to memory issues, so we always create those in the serial part of the code and then simply fill them in parallel.
+
+
+SIMD
+^^^^
+For simple computations that are compute bound, we use SIMD (`Single Instruction Multiple Data <https://en.wikipedia.org/wiki/Single_instruction,_multiple_data>`_) instructions to make use of the AVX/AVX2/AVX512 instruction sets on modern CPUs. To simplify the use of these instructions, we use the `xsimd <https://github.com/xtensor-stack/xsimd>`_ library.
+
+CMake
+^^^^^
+
+When editing the C++ code, it may be useful to use ``CMake`` and ``make`` directly to only recompile those parts of the code that changed. This can be achieved as below
+
+.. code-block::
+
+    git clone --recursive git@github.com:hiddenSymmetries/simsopt.git
+    cd simsopt
+    pip3 install -e .
+    mkdir cmake-build
+    cd cmake-build
+    cmake ..
+    make -j
+    cd ../src
+    rm simsoptpp.cpython-38-x86_64-linux-gnu.so
+    ln -s ../cmake-build/simsoptpp.cpython-38-x86_64-linux-gnu.so .
+
+You may have to adjust the last two lines to match your local system.
+From then on, you can always just call ``make -j`` inside the ``cmake-build`` directory to recompile the C++ part of the code.

docs/source/index.rst

@@ -74,6 +74,7 @@ optimization.  Others include `STELLOPT
    source
    publications
    contributing
+   cpp
 
 .. toctree::
    :maxdepth: 3

examples/1_Simple/tracing_fieldline.py

@@ -1,5 +1,20 @@
 #!/usr/bin/env python3
-# import matplotlib; matplotlib.use('agg')  # noqa
+
+"""
+This example demonstrates how to use SIMSOPT to compute Poincare plots.
+"""
+
+import time
+import os
+import logging
+import sys
+import numpy as np
+try:
+    from mpi4py import MPI
+    comm = MPI.COMM_WORLD
+except ImportError:
+    comm = None
+
 from simsopt.field.biotsavart import BiotSavart
 from simsopt.field.magneticfieldclasses import InterpolatedField, UniformInterpolationRule
 from simsopt.geo.surfacexyztensorfourier import SurfaceRZFourier
@@ -8,26 +23,15 @@
     particles_to_vtk, compute_fieldlines, LevelsetStoppingCriterion, plot_poincare_data
 from simsopt.geo.curve import curves_to_vtk
 from simsopt.util.zoo import get_ncsx_data
-import simsoptpp as sopp
-try:
-    from mpi4py import MPI
-    comm = MPI.COMM_WORLD
-except ImportError:
-    comm = None
 
-import numpy as np
-import time
-import os
-import logging
-import sys
+print("Running 1_Simple/tracing_fieldline.py")
+print("=====================================")
+
 sys.path.append(os.path.join("..", "tests", "geo"))
 logging.basicConfig()
 logger = logging.getLogger('simsopt.field.tracing')
 logger.setLevel(1)
 
-print("Running 1_Simple/tracing_fieldline.py")
-print("=====================================")
-
 # check whether we're in CI, in that case we make the run a bit cheaper
 ci = "CI" in os.environ and os.environ['CI'].lower() in ['1', 'true']
 nfieldlines = 3 if ci else 30
@@ -38,14 +42,10 @@
 OUT_DIR = "./output/"
 os.makedirs(OUT_DIR, exist_ok=True)
 
-"""
-This examples demonstrate how to use SIMSOPT to compute Poincare plots and
-guiding center trajectories of particles
-"""
-
 
+nfp = 3
 curves, currents, ma = get_ncsx_data()
-coils = coils_via_symmetries(curves, currents, 3, True)
+coils = coils_via_symmetries(curves, currents, nfp, True)
 curves = [c.curve for c in coils]
 bs = BiotSavart(coils)
 print("Mean(|B|) on axis =", np.mean(np.linalg.norm(bs.set_points(ma.gamma()).B(), axis=1)))
@@ -55,15 +55,12 @@
 mpol = 5
 ntor = 5
 stellsym = True
-nfp = 3
-phis = np.linspace(0, 1, nfp*2*ntor+1, endpoint=False)
-thetas = np.linspace(0, 1, 2*mpol+1, endpoint=False)
-s = SurfaceRZFourier(
-    mpol=mpol, ntor=ntor, stellsym=stellsym, nfp=nfp, quadpoints_phi=phis, quadpoints_theta=thetas)
+s = SurfaceRZFourier(mpol=mpol, ntor=ntor, stellsym=stellsym, nfp=nfp,
+                     range="full torus", nphi=64, ntheta=24)
 s.fit_to_curve(ma, 0.70, flip_theta=False)
 
 s.to_vtk(OUT_DIR + 'surface')
-sc_fieldline = SurfaceClassifier(s, h=0.1, p=2)
+sc_fieldline = SurfaceClassifier(s, h=0.03, p=2)
 sc_fieldline.to_vtk(OUT_DIR + 'levelset', h=0.02)
 
 
@@ -86,16 +83,43 @@ def trace_fieldlines(bfield, label):
 # trace_fieldlines(bs, 'bs')
 
 
-n = 16
+# Bounds for the interpolated magnetic field chosen so that the surface is
+# entirely contained in it
+n = 20
 rs = np.linalg.norm(s.gamma()[:, :, 0:2], axis=2)
 zs = s.gamma()[:, :, 2]
 rrange = (np.min(rs), np.max(rs), n)
-phirange = (0, 2*np.pi/3, n*2)
+phirange = (0, 2*np.pi/nfp, n*2)
+# exploit stellarator symmetry and only consider positive z values:
 zrange = (0, np.max(zs), n//2)
+
+
+def skip(rs, phis, zs):
+    # The RegularGrindInterpolant3D class allows us to specify a function that
+    # is used in order to figure out which cells to be skipped.  Internally,
+    # the class will evaluate this function on the nodes of the regular mesh,
+    # and if *all* of the eight corners are outside the domain, then the cell
+    # is skipped.  Since the surface may be curved in a way that for some
+    # cells, all mesh nodes are outside the surface, but the surface still
+    # intersects with a cell, we need to have a bit of buffer in the signed
+    # distance (essentially blowing up the surface a bit), to avoid ignoring
+    # cells that shouldn't be ignored
+    rphiz = np.asarray([rs, phis, zs]).T.copy()
+    dists = sc_fieldline.evaluate_rphiz(rphiz)
+    skip = list((dists < -0.05).flatten())
+    print("Skip", sum(skip), "cells out of", len(skip), flush=True)
+    return skip
+
+
 bsh = InterpolatedField(
-    bs, degree, rrange, phirange, zrange, True, nfp=3, stellsym=True
+    bs, degree, rrange, phirange, zrange, True, nfp=nfp, stellsym=True, skip=skip
 )
-# print('Error in B', bsh.estimate_error_B(1000), flush=True)
+
+bsh.set_points(ma.gamma().reshape((-1, 3)))
+bs.set_points(ma.gamma().reshape((-1, 3)))
+Bh = bsh.B()
+B = bs.B()
+print("|B-Bh| on axis", np.sort(np.abs(B-Bh).flatten()))
 trace_fieldlines(bsh, 'bsh')
 print("End of 1_Simple/tracing_fieldline.py")
 print("=====================================")

examples/1_Simple/tracing_particle.py

@@ -1,4 +1,21 @@
 #!/usr/bin/env python3
+
+"""
+This examples demonstrate how to use SIMSOPT to compute guiding center
+trajectories of particles in cylindrical coordinates for a vacuum field.
+"""
+
+import time
+import os
+import logging
+import sys
+import numpy as np
+try:
+    from mpi4py import MPI
+    comm = MPI.COMM_WORLD
+except ImportError:
+    comm = None
+
 from simsopt.field.biotsavart import BiotSavart
 from simsopt.field.magneticfieldclasses import InterpolatedField, UniformInterpolationRule
 from simsopt.geo.surfacexyztensorfourier import SurfaceRZFourier
@@ -8,25 +25,15 @@
 from simsopt.geo.curve import curves_to_vtk
 from simsopt.util.zoo import get_ncsx_data
 from simsopt.util.constants import PROTON_MASS, ELEMENTARY_CHARGE, ONE_EV
-import simsoptpp as sopp
-try:
-    from mpi4py import MPI
-    comm = MPI.COMM_WORLD
-except ImportError:
-    comm = None
 
-import numpy as np
-import time
-import os
-import logging
-import sys
+print("Running 1_Simple/tracing_particle.py")
+print("====================================")
+
 sys.path.append(os.path.join("..", "tests", "geo"))
 logging.basicConfig()
 logger = logging.getLogger('simsopt.field.tracing')
 logger.setLevel(1)
 
-print("Running 1_Simple/tracing_particle.py")
-print("====================================")
 # check whether we're in CI, in that case we make the run a bit cheaper
 ci = "CI" in os.environ and os.environ['CI'].lower() in ['1', 'true']
 nparticles = 3 if ci else 100
@@ -36,14 +43,9 @@
 OUT_DIR = "./output/"
 os.makedirs(OUT_DIR, exist_ok=True)
 
-"""
-This examples demonstrate how to use SIMSOPT to compute guiding center
-trajectories of particles
-"""
-
-
+nfp = 3
 curves, currents, ma = get_ncsx_data()
-coils = coils_via_symmetries(curves, currents, 3, True)
+coils = coils_via_symmetries(curves, currents, nfp, True)
 curves = [c.curve for c in coils]
 bs = BiotSavart(coils)
 print("Mean(|B|) on axis =", np.mean(np.linalg.norm(bs.set_points(ma.gamma()).B(), axis=1)))
@@ -53,11 +55,8 @@
 mpol = 5
 ntor = 5
 stellsym = True
-nfp = 3
-phis = np.linspace(0, 1, nfp*2*ntor+1, endpoint=False)
-thetas = np.linspace(0, 1, 2*mpol+1, endpoint=False)
-s = SurfaceRZFourier(
-    mpol=mpol, ntor=ntor, stellsym=stellsym, nfp=nfp, quadpoints_phi=phis, quadpoints_theta=thetas)
+s = SurfaceRZFourier(mpol=mpol, ntor=ntor, stellsym=stellsym, nfp=nfp,
+                     range="full torus", nphi=64, ntheta=24)
 
 
 s.fit_to_curve(ma, 0.20, flip_theta=False)
@@ -67,10 +66,11 @@
 zs = s.gamma()[:, :, 2]
 
 rrange = (np.min(rs), np.max(rs), n)
-phirange = (0, 2*np.pi/3, n*2)
+phirange = (0, 2*np.pi/nfp, n*2)
+# exploit stellarator symmetry and only consider positive z values:
 zrange = (0, np.max(zs), n//2)
 bsh = InterpolatedField(
-    bs, degree, rrange, phirange, zrange, True, nfp=3, stellsym=True
+    bs, degree, rrange, phirange, zrange, True, nfp=nfp, stellsym=True
 )
 
 

pyproject.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["setuptools>=45", "wheel", "numpy", "cmake", "ninja", "setuptools_scm[toml]>=6.0"]
+requires = ["setuptools>=45", "wheel", "oldest-supported-numpy", "cmake", "ninja", "setuptools_scm[toml]>=6.0"]
 build-backend = "setuptools.build_meta"
 
 [tools.setuptools_scm]

src/simsopt/field/magneticfieldclasses.py

@@ -517,7 +517,7 @@ class InterpolatedField(sopp.InterpolatedField, MagneticField):
     This resulting interpolant can then be evaluated very quickly.
     """
 
-    def __init__(self, field, degree, rrange, phirange, zrange, extrapolate=True, nfp=1, stellsym=False):
+    def __init__(self, field, degree, rrange, phirange, zrange, extrapolate=True, nfp=1, stellsym=False, skip=None):
         r"""
         Args:
             field: the underlying :mod:`simsopt.field.magneticfield.MagneticField` to be interpolated.
@@ -535,17 +535,34 @@ def __init__(self, field, degree, rrange, phirange, zrange, extrapolate=True, nf
             stellsym: Whether to exploit stellarator symmetry. In this case
                       ``z`` is always mapped to be positive, hence it makes sense to use
                       ``zmin=0``.
+            skip: a function that takes in a point (in cylindrical (r,phi,z)
+                  coordinates) and returns whether to skip that location when
+                  building the interpolant or not. The signature should be
+
+                  .. code-block:: Python
+
+                      def skip(r: double, phi: double, z: double) -> bool:
+                          ...
+
+                  See also here
+                  https://github.com/hiddenSymmetries/simsopt/pull/227 for a
+                  graphical illustration.
+
         """
         MagneticField.__init__(self)
         if stellsym and zrange[0] != 0:
             logger.warning(fr"Sure about zrange[0]={zrange[0]}? When exploiting stellarator symmetry, the interpolant is never evaluated for z<0.")
         if nfp > 1 and abs(phirange[1] - 2*np.pi/nfp) > 1e-14:
             logger.warning(fr"Sure about phirange[1]={phirange[1]}? When exploiting rotational symmetry, the interpolant is never evaluated for phi>2\pi/nfp.")
 
-        sopp.InterpolatedField.__init__(self, field, degree, rrange, phirange, zrange, extrapolate, nfp, stellsym)
+        if skip is None:
+            def skip(xs, ys, zs):
+                return [False for _ in xs]
+
+        sopp.InterpolatedField.__init__(self, field, degree, rrange, phirange, zrange, extrapolate, nfp, stellsym, skip)
         self.__field = field
 
-    def to_vtk(self, filename, h=0.1):
+    def to_vtk(self, filename):
         """Export the field evaluated on a regular grid for visualisation with e.g. Paraview."""
         degree = self.rule.degree
         MagneticField.to_vtk(