From 4bb2d42112f18bfe30d45da9fda97701040137d2 Mon Sep 17 00:00:00 2001 From: Antoine Baillod Date: Mon, 20 May 2024 14:18:00 -0400 Subject: [PATCH] Renamed the function create_equally_spaced_windowpane_curves, updated an example --- .../optimize_coil_position_orientation.py | 11 +-- src/simsopt/geo/curve.py | 4 +- src/simsopt/geo/curveobjectives.py | 68 ++++++++++++++++++- 3 files changed, 75 insertions(+), 8 deletions(-) diff --git a/examples/1_Simple/optimize_coil_position_orientation.py b/examples/1_Simple/optimize_coil_position_orientation.py index dc6e9d37c..796b32b71 100644 --- a/examples/1_Simple/optimize_coil_position_orientation.py +++ b/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 diff --git a/src/simsopt/geo/curve.py b/src/simsopt/geo/curve.py index 40f68b950..dcc726736 100644 --- a/src/simsopt/geo/curve.py +++ b/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 diff --git a/src/simsopt/geo/curveobjectives.py b/src/simsopt/geo/curveobjectives.py index 14a0c68e2..9f2082409 100644 --- a/src/simsopt/geo/curveobjectives.py +++ b/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 )