Renamed the function create_equally_spaced_windowpane_curves, updated…

… an example

examples/1_Simple/optimize_coil_position_orientation.py

@@ -19,7 +19,7 @@
 import numpy as np
 from scipy.optimize import minimize
 
-from simsopt.geo import SurfaceRZFourier, create_equally_spaced_windowpane_curves, \
+from simsopt.geo import SurfaceRZFourier, create_equally_spaced_oriented_curves, \
     CurveLength, curves_to_vtk, create_equally_spaced_curves
 from simsopt.field import Current, coils_via_symmetries, BiotSavart
 from simsopt.objectives import SquaredFlux
@@ -58,7 +58,7 @@
 
 # Create the initial coils:
 base_tf_curves = create_equally_spaced_curves(n_tf_coils, s.nfp, stellsym=True, R0=R0, R1=R1, order=2)
-base_wp_curves = create_equally_spaced_windowpane_curves(n_wp_coils, s.nfp, True, R0=(R0+R1)*1.01, R1=R1/10, Z0=0, order=2)
+base_wp_curves = create_equally_spaced_oriented_curves(n_wp_coils, s.nfp, R0=(R0+R1)*1.01, R1=R1/10, Z0=0, order=2)
 
 # We scale the currents so that dofs have all the same order of magnitude
 base_tf_currents = [ScaledCurrent(Current(1.0), 1e5) for i in range(n_tf_coils)]
@@ -71,10 +71,10 @@
 # We also fix the wp coils geometry, but keep their position and orientation unfixed
 for c in base_wp_curves:
     c.fix_all()
-    for xyz in ['x','y','z']:
-        c.unfix(f'{xyz}0')
+    for xyz in ['x0','y0','z0']:
+        c.unfix( xyz )
     for ypr in ['yaw', 'pitch', 'roll']:
-        c.unfix(f'{ypr}')
+        c.unfix( ypr )
 
 # We unfix all currents
 for c in base_tf_currents:
@@ -87,6 +87,7 @@
 # of the currents:
 base_tf_currents[0].fix_all()
 
+# Generate full array of coils using symmetries
 tf_coils = coils_via_symmetries(base_tf_curves, base_tf_currents, s.nfp, True)
 wp_coils = coils_via_symmetries(base_wp_curves, base_wp_currents, s.nfp, True)
 coils = tf_coils + wp_coils

src/simsopt/geo/curve.py

@@ -12,7 +12,7 @@
 from .._core.derivative import derivative_dec
 from .plotting import fix_matplotlib_3d
 
-__all__ = ['Curve', 'JaxCurve', 'RotatedCurve', 'curves_to_vtk', 'create_equally_spaced_curves', 'create_equally_spaced_windowpane_curves', 'Curve2D', 'CurveCWSFourier', 'create_equally_spaced_planar_curves']
+__all__ = ['Curve', 'JaxCurve', 'RotatedCurve', 'curves_to_vtk', 'create_equally_spaced_curves', 'create_equally_spaced_oriented_curves', 'Curve2D', 'CurveCWSFourier', 'create_equally_spaced_planar_curves']
 
 @jit
 def incremental_arclength_pure(d1gamma):
@@ -884,7 +884,7 @@ def create_equally_spaced_curves(ncurves, nfp, stellsym, R0=1.0, R1=0.5, order=6
         curves.append(curve)
     return curves
 
-def create_equally_spaced_windowpane_curves( ncurves, nfp, R0, R1, Z0, order, numquadpoints=None ):
+def create_equally_spaced_oriented_curves( ncurves, nfp, R0, R1, Z0, order, numquadpoints=None ):
     if numquadpoints is None:
         numquadpoints = 15 * order
 

src/simsopt/geo/curveobjectives.py

@@ -15,7 +15,7 @@
 
 __all__ = ['CurveLength', 'LpCurveCurvature', 'LpCurveTorsion',
            'CurveCurveDistance', 'CurveSurfaceDistance', 'ArclengthVariation',
-           'MeanSquaredCurvature', 'LinkingNumber', 'CurveCylinderDistance', 'FramedCurveTwist']
+           'MeanSquaredCurvature', 'LinkingNumber', 'CurveCylinderDistance', 'FramedCurveTwist', 'MinCurveCurveDistance']
 
 
 @jit
@@ -721,3 +721,69 @@ def dJ(self):
             grad += self.framedcurve.curve.dgammadash_by_dcoeff_vjp(grad1)
 
         return grad 
+
+
+
+def max_distance_pure(g1, g2, dmax, p):
+    """
+    This returns 0 if all points of g1 have at least one point of g2 at a distance smaller or equal to dmax
+    Otherwise, returns the sum of |g2-g1_i|-dmax where only points further than dmax are considered.
+    The minimum distance between a point g1_i and g2 is obtained using the p-norm, with p < -1.
+    """
+    dists = jnp.sqrt(jnp.sum( (g1[:, None, :] - g2[None, :, :])**2, axis=2))
+
+    # Estimate min of dists using p-norm. The minimum is taken along the axis=1. mindists is then an array of length g1.size, where mindists[i]=min_j(|g1[i]-g2[j]|)
+    mindists = jnp.sum(dists**p, axis=1)**(1./p)
+
+
+    # We now evaluate if any of mindists is larger than dmax. If yes, we add the value of (mindists[i]-dmax)**2 to the output. 
+    # We normalize by the number of quadrature points along the first curve g1.
+    return jnp.sum(jnp.maximum(mindists-dmax, 0)**2) / g1.shape[0]
+
+
+class MinCurveCurveDistance(Optimizable):
+    """
+    This class can be used to constrain a curve to remain close
+    to another curve.
+    """
+    def __init__(self, curve1, curve2, maximum_distance, p=-10):
+        self.curve1 = curve1
+        self.curve2 = curve2
+        self.maximum_distance = maximum_distance
+
+        self.J_jax = lambda g1, g2: max_distance_pure(g1, g2, self.maximum_distance, p)
+        self.this_grad_0 = jit(lambda g1, g2: grad(self.J_jax, argnums=0)(g1, g2))
+        self.this_grad_1 = jit(lambda g1, g2: grad(self.J_jax, argnums=1)(g1, g2))
+
+        Optimizable.__init__(self, depends_on=[curve1, curve2])
+
+    def max_distance(self):
+        """
+        returns the distance the most isolated point of curve1 to curve2
+        """
+        g1 = self.curve1.gamma()
+        g2 = self.curve2.gamma()
+
+        # Evaluate all distances
+        dists = np.sqrt(np.sum( (g1[:, None, :] - g2[None, :, :])**2, axis=2))
+
+        # Find all min distances
+        mindists = np.min(dists, axis=1)
+
+        return np.max(mindists)
+
+    def J(self):
+        g1 = self.curve1.gamma()
+        g2 = self.curve2.gamma()
+
+        return self.J_jax( g1, g2 )
+
+    @derivative_dec
+    def dJ(self):
+        g1 = self.curve1.gamma()
+        g2 = self.curve2.gamma()
+
+        grad0 = self.this_grad_0(g1, g2)
+        grad1 = self.this_grad_1(g1, g2)
+
+        return self.curve1.dgamma_by_dcoeff_vjp( grad0 ) + self.curve2.dgamma_by_dcoeff_vjp( grad1 )