quadcoil 0.1.0__py3-none-any.whl

quadcoil/__init__.py

@@ -0,0 +1,13 @@
+from .math_utils import *
+from .quadcoil_params import *
+from .surfacerzfourier_jax import *
+from .winding_surface import *
+from .wrapper import *
+from .solver import *
+from .quadcoil import *
+from . import io
+from . import quantity
+# from .conf import *
+# All submodules uses quadcoil. So, we will
+# not import them here to avoid circular 
+# imports.

quadcoil/io/__init__.py

@@ -0,0 +1,7 @@
+from .jax import *
+from .file import *
+from .desc import *
+from .simsopt import *
+from .coil_cutting import *
+from .focus import *
+

quadcoil/io/coil_cutting.py

@@ -0,0 +1,257 @@
+
+# Import packages.
+import numpy as np
+import matplotlib.pyplot as plt
+from scipy.fft import rfft
+# Importing simsopt 
+from quadcoil import QuadcoilParams
+from quadcoil.quantity import Phi_with_net_current
+
+def coil_zeta_theta_from_qp(
+    qp:QuadcoilParams,
+    dofs,
+    coils_per_half_period=5,
+    theta_shift=0): 
+    stellsym = qp.stellsym
+    nzeta_coil = len(qp.quadpoints_phi)
+    nfp = qp.nfp 
+    theta = qp.quadpoints_theta * 2 * np.pi 
+    zeta = qp.quadpoints_phi * 2 * np.pi 
+    net_poloidal_current_amperes = qp.net_poloidal_current_amperes 
+
+    # ------------------------
+    # Load current potential
+    # ------------------------
+    current_potential = Phi_with_net_current(qp, dofs)
+
+    # Now just generate a new monotonic array with the correct first value:
+    theta = theta[0] + np.linspace(0,2*np.pi,len(theta),endpoint=False)
+
+    d = 2*np.pi/nfp
+    # We add a field period to the "left" and "right" of the 
+    # present field period so that no contours are cutoff.
+    zeta_3 = np.concatenate((zeta-d, zeta, zeta+d))
+    # If the net current is non-neglegible, use the net current to calculate spacing.
+    # Else, use the maximum phi to calculate spacing.
+    # we normalize the current potential correspondingly into "data",
+    # and also pad it to contain 3 field periods.
+    if abs(net_poloidal_current_amperes) > np.finfo(float).eps:
+        data = current_potential / net_poloidal_current_amperes * nfp  
+        data_3 = np.concatenate((data-1,data,data+1))
+    else:
+        data = current_potential / np.max(current_potential)
+        data_3 = np.concatenate((data,data,data))
+
+    # We only generate contours based on the "middle"
+    # field period. 
+    d = 0.5/coils_per_half_period
+    level_shift = theta_shift % 1 * d
+    if stellsym:
+        if theta_shift != 0:
+            print('Warning: non-zero theta-shift is not permitted when stell-sym is true.')
+        levels = np.linspace(0,0.5,coils_per_half_period,endpoint=False)
+    else:
+        levels = np.linspace(0,1,coils_per_half_period*2,endpoint=False) + level_shift
+    # The contours are symmetric w.r.t. 2pi/nfp/2, i.e. stellarator
+    # symmetric, by default.
+    levels = levels + d/2
+    fig = plt.figure()
+    cdata = plt.contour(zeta_3, theta, np.transpose(data_3), levels, colors='k')
+    plt.close(fig)
+    num_coils_found = len(cdata.levels)
+    if num_coils_found != 2*coils_per_half_period:
+        print("Warning: The expected number of levels != 2 * coils_per_half_period.")
+
+    contour_zeta=[]
+    contour_theta=[]
+    num_coils = 0
+    for j in range(num_coils_found):
+        paths_in_level = cdata.allsegs[j]
+        # Each level can contain many paths.
+        for p in paths_in_level:
+            plt.plot(p[:, 0], p[:, 1])
+            # These are the zeta (toroidal angles)
+            # and theta of the paths.
+            contour_zeta.append(p[:, 0])
+            contour_theta.append(p[:, 1])
+            num_coils += 1
+
+    print('Cutting complete. Number of coils/half field period:', num_coils)
+    return(contour_zeta, contour_theta)
+
+# IFFT a array in real space to a sin/cos series used by sinsopt.geo.curve
+def ifft_simsopt(x, order):
+    assert len(x) >= 2*order  # the order of the fft is limited by the number of samples
+    xf = rfft(x) / len(x)
+
+    fft_0 = [xf[0].real]  # find the 0 order coefficient
+    fft_cos = 2 * xf[1:order + 1].real  # find the cosine coefficients
+    fft_sin = (-2 * xf[:order + 1].imag)[1:]  # find the sine coefficients
+    dof = np.zeros(order*2+1)
+    dof[0] = fft_0[0]
+    dof[1::2] = fft_sin
+    dof[2::2] = fft_cos
+
+    return dof
+
+# This script assumes the contours do not zig-zag back and forth across the theta=0 line,
+# after shifting the current potential by theta_shift grid points.
+# The nescin file is used to provide the coil winding surface, so make sure this is consistent with the regcoil run.
+# ilambda is the index in the lambda scan which you want to select.
+# def cut_coil(qp, qpst):
+# filename = 'regcoil_out.li383.nc' # sys.argv[1]
+# TODO: these tow have a lot of duplicate code. Clean up when finalizing how QUADCOIL is integrated.
+def coil_xyz_from_qp(
+    qp:QuadcoilParams,
+    dofs,
+    coils_per_half_period=1,
+    theta_shift=0,
+    save=False, save_name='placeholder'): 
+    contour_zeta, contour_theta = coil_zeta_theta_from_qp(
+        qp=qp,
+        dofs=dofs,
+        coils_per_half_period=coils_per_half_period,
+        theta_shift=theta_shift
+    )
+    num_coils = len(contour_zeta)    
+    nfp = qp.nfp 
+    net_poloidal_current_amperes = qp.net_poloidal_current_amperes 
+
+    # ------------------------
+    # Load surface shape
+    # ------------------------
+    contour_R = []
+    contour_Z = []
+    for j in range(num_coils):
+        contour_R.append(np.zeros_like(contour_zeta[j]))
+        contour_Z.append(np.zeros_like(contour_zeta[j]))
+
+    surf = qp.winding_surface.to_simsopt()
+    for m in range(surf.mpol+1): # 0 to mpol
+        for i in range(2*surf.ntor+1):
+            n = i-surf.ntor
+            crc = surf.rc[m, i]
+            czs = surf.zs[m, i]
+            if surf.stellsym:
+                crs = 0
+                czc = 0
+            else: # Returns ValueError for stellsym cases
+                crs = surf.get_rs(m, n)
+                czc = surf.get_zc(m, n)
+            for j in range(num_coils):
+                angle = m*contour_theta[j] - n*contour_zeta[j]*surf.nfp
+                # Was filled with zeroes.
+                # Are lists because contou lengths are not uniform.
+                contour_R[j] = contour_R[j] + crc*np.cos(angle) + crs*np.sin(angle)
+                contour_Z[j] = contour_Z[j] + czs*np.sin(angle) + czc*np.cos(angle)
+
+    contour_X = []
+    contour_Y = []
+    for j in range(num_coils):
+        contour_X.append(contour_R[j]*np.cos(contour_zeta[j]))
+        contour_Y.append(contour_R[j]*np.sin(contour_zeta[j]))
+    coil_currents = net_poloidal_current_amperes / num_coils / qp.nfp
+    if qp.stellsym:
+        coil_currents = coil_currents/2
+    if save: 
+        coilsFilename='coils.'+save_name
+        print("coilsFilename:",coilsFilename)
+        # Write coils file
+        f = open(coilsFilename,'w')
+        f.write('periods '+str(nfp)+'\n')
+        f.write('begin filament\n')
+        f.write('mirror NIL\n')
+
+        for j in range(num_coils):
+            N = len(contour_X[j])
+            for k in range(N):
+                f.write('{:14.22e} {:14.22e} {:14.22e} {:14.22e}\n'.format(contour_X[j][k],contour_Y[j][k],contour_Z[j][k],coil_currents))
+            # Close the loop
+            k=0
+            f.write('{:14.22e} {:14.22e} {:14.22e} {:14.22e} 1 Modular\n'.format(contour_X[j][k],contour_Y[j][k],contour_Z[j][k],0))
+
+        f.write('end\n')
+        f.close()
+    return (
+        contour_X,
+        contour_Y,
+        contour_Z,
+        coil_currents,
+        # np.sqrt(minSeparation2)
+    )
+
+# Load curves from lists of arrays containing x, y, and z.
+def simsopt_curves_from_xyz(
+    contour_X,
+    contour_Y,
+    contour_Z, 
+    order=None, ppp=20):
+    num_coils = len(contour_X)
+    try:     
+        from simsopt.geo import CurveXYZFourier
+    except:
+        raise ImportError('Simsopt is required to use the coil-cutting features.')
+    # Calculating order
+    if not order:
+        order=float('inf')
+        for i in range(num_coils):
+            xArr = contour_X[i]
+            yArr = contour_Y[i]
+            zArr = contour_Z[i]
+            for x in [xArr, yArr, zArr]:
+                if len(x)//2<order:
+                    order = len(x)//2
+    
+    coils = [CurveXYZFourier(order*ppp, order) for i in range(num_coils)]
+    # Compute the Fourier coefficients for each coil
+    for ic in range(num_coils):
+        xArr = contour_X[ic]
+        yArr = contour_Y[ic]
+        zArr = contour_Z[ic]
+
+        # Compute the Fourier coefficients
+        dofs=[]
+        for x in [xArr, yArr, zArr]:
+            dof_i = ifft_simsopt(x, order)
+            dofs.append(dof_i)
+
+        coils[ic].local_x = np.concatenate(dofs)
+    return coils
+
+def simsopt_coil_from_qp(
+    qp, dofs, coils_per_half_period, theta_shift=0,
+    method=coil_xyz_from_qp,
+    base_mode=False,
+    order=10, ppp=40):
+    try:     
+        from simsopt.field import Current # , Coil
+        from simsopt.field.coil import coils_via_symmetries
+    except:
+        raise ImportError('Simsopt is required to use the coil-cutting features.')
+    (
+        contour_X,
+        contour_Y,
+        contour_Z,
+        coil_currents,
+        # min_separation
+    ) = method(
+        qp=qp,
+        dofs=dofs,
+        coils_per_half_period=coils_per_half_period,
+        theta_shift=theta_shift,
+        save=False
+    )
+    curves = simsopt_curves_from_xyz(
+        contour_X,
+        contour_Y,
+        contour_Z, 
+        order=order,
+        ppp=ppp)
+    currents=[]
+    for i in range(len(curves)):
+        currents.append(Current(coil_currents))
+    if base_mode:
+        return curves, currents
+    coils = coils_via_symmetries(curves, currents, qp.nfp, qp.stellsym)
+    return coils
+

quadcoil/io/desc.py

@@ -0,0 +1,195 @@
+from quadcoil import quadcoil, SurfaceRZFourierJAX, get_quantity
+import jax.numpy as jnp
+from scipy.constants import mu_0
+
+
+def generate_desc_scaling(objective_name, constraint_name, scales):
+    '''
+    A helper method for calculating the units of each objectives and constraints
+    using quantities scales = desc.objectives.normalization.compute_scaling_factors(eq).
+    The formula must be first implemented as quadcoil.objective.<quantity_name>_desc_unit.
+    '''
+    if isinstance(objective_name, str):
+        objective_unit = get_quantity(objective_name).desc_unit(scales)
+    else:
+        objective_unit_new = []
+        for obj_name in objective_name:
+            objective_unit_new.append(get_quantity(obj_name).desc_unit(scales))
+        objective_unit = tuple(objective_unit_new)
+    
+    constraint_unit_new = []
+    for cons_name in constraint_name:
+        constraint_unit_new.append(get_quantity(cons_name).desc_unit(scales))
+    constraint_unit = tuple(constraint_unit_new)
+    return(objective_unit, constraint_unit)
+
+# ----- An earlier, simpler wrapper. May be discontinued later. -----
+def quadcoil_desc(
+    desc_eq,
+    vacuum:bool,
+    plasma_M_theta:int,
+    plasma_N_phi:int,
+    desc_scaling:bool=True,
+    **kwargs
+):
+    '''
+    A simple DESC interface for ``quadcoil.quadcoil``, that replaces the following
+    parameters using information extracted from DESC:
+    
+    .. code-block:: python
+
+        nfp
+        stellsym
+        plasma_mpol
+        plasma_ntor
+        plasma_quadpoints_phi
+        plasma_quadpoints_theta
+        plasma_dofs
+        net_poloidal_current_amperes
+        Bnormal_plasma
+        winding_dofs
+
+    '''
+    try:
+        from desc.integrals import compute_B_plasma
+        from desc.grid import LinearGrid
+        from desc.objectives.normalization import compute_scaling_factors
+    except:
+        raise ModuleNotFoundError('DESC must be installed for the... DESC interface to run.')
+    if 'plasma_coil_distance' in kwargs:
+        # DESC surfaces have different handed-ness 
+        # compared to simsopt surfaces. Because QUADCOIL
+        # uses simsopt conventions, we flip the sign 
+        # of offset direction here.
+        kwargs['plasma_coil_distance'] = -kwargs['plasma_coil_distance'] 
+    # We want to make sure DESC and QUADCOIL
+    # uses the same grid on the plasma surface.
+    # In case we want to use Bnormal_plasma 
+    # from DESC.
+    # To do this, we create a grid inside DESC.
+    surface_grid = LinearGrid(
+        NFP=desc_eq.NFP,
+        # If we set this to sym it'll only evaluate 
+        # theta from 0 to pi.
+        sym=False, 
+        M=plasma_M_theta,#Poloidal grid resolution.
+        N=plasma_N_phi,
+        rho=1.0
+    )
+    # Recovering quadpoints from the LinearGrid
+    # for use in QUADCOIL and simsopt
+    desc_plasma_quadpoints_theta = surface_grid.nodes[surface_grid.unique_theta_idx, 1]/jnp.pi/2
+    desc_plasma_quadpoints_phi = surface_grid.nodes[surface_grid.unique_zeta_idx, 2]/jnp.pi/2
+    desc_plasma_surf = SurfaceRZFourierJAX.from_desc(
+        desc_eq.surface, 
+        desc_plasma_quadpoints_phi,
+        desc_plasma_quadpoints_theta
+    )
+    desc_net_poloidal_current_amperes = (
+        -desc_eq.compute("G", grid=LinearGrid(rho=jnp.array(1.0)))["G"][0]
+        / mu_0
+        * 2
+        * jnp.pi
+    )
+    # Calculating B_plasma if not vacuum
+    if not vacuum:
+        desc_Bnormal_plasma = compute_B_plasma(
+            desc_eq, 
+            eval_grid=surface_grid, 
+            normal_only=True
+            # chunk_size=B_plasma_chunk_size
+        )
+        # Reshape the output so that axis=0 is zeta (the toroidal coordinate)
+        desc_Bnormal_plasma = surface_grid.meshgrid_reshape(desc_Bnormal_plasma, order='rzt')[0]
+    else:
+        desc_Bnormal_plasma = jnp.zeros((
+            len(desc_plasma_quadpoints_phi),
+            len(desc_plasma_quadpoints_theta)
+        ))
+    print('Finished reading DESC equilibrium. G =', desc_net_poloidal_current_amperes)
+    print('Finished reading DESC equilibrium. G =', type(desc_net_poloidal_current_amperes))
+    # Detect if the user has provided any arguments 
+    # that will also-be auto-calculated using DESC. 
+    # if any, these objectives will be discarded.
+    redundant_arg_names = set(_DESC_DERIVED_ARGNAMES) & kwargs.keys()
+    if redundant_arg_names:
+        warnings.warn(
+            'Redundant arguments detected: ' 
+            + str(redundant_arg_names)
+            + '. These arguments must be calculated from '\
+            'the equilibrium. The provided values will be discarded.'
+        )
+    # Add all the equilibrium-dependent, non-differentiable args to quadcopil_args
+    # Non-differentiable 
+    # (This distinction is only for internal reference.
+    # the desc _Objective wrapper of QUADCOIL uses a custom_jvp rule in 
+    # quadcoil.io.jax, and that function makes the distinction between
+    # differentiable and non-differentiable args.)
+    kwargs['nfp'] = desc_eq.NFP
+    kwargs['stellsym'] = desc_eq.sym
+    kwargs['plasma_mpol'] = desc_plasma_surf.mpol
+    kwargs['plasma_ntor'] = desc_plasma_surf.ntor
+    kwargs['plasma_quadpoints_phi'] = desc_plasma_quadpoints_phi
+    kwargs['plasma_quadpoints_theta'] = desc_plasma_quadpoints_theta
+    # Differentiable 
+    kwargs['plasma_dofs'] = desc_plasma_surf.dofs
+    kwargs['net_poloidal_current_amperes'] = desc_net_poloidal_current_amperes
+    kwargs['Bnormal_plasma'] = desc_Bnormal_plasma
+    # For unconstrained cases, kwargs will not contain 
+    # 'constraint_name'. But we still want to use the scaling helper function.
+    # So we manually add a blank entry.
+    if 'constraint_name' not in kwargs:
+        kwargs['constraint_name'] = ()
+    if desc_scaling:
+        print('Overriding provided/default scaling with DESC scaling.')
+        obj_unit_new, cons_unit_new = generate_desc_scaling(
+            kwargs['objective_name'], 
+            kwargs['constraint_name'], 
+            compute_scaling_factors(desc_eq)
+        )
+        kwargs['objective_unit'] = obj_unit_new
+        kwargs['constraint_unit'] = cons_unit_new
+    return(
+        quadcoil(**kwargs)
+    )
+
+def load_FourierCurrentPotentialField(eq, scf, winding_quadpoints_phi=None, winding_quadpoints_theta=None, source_grid=None):
+    try:
+        from desc.vmec_utils import ptolemy_identity_rev
+    except:
+        raise ModuleNotFoundError(
+            'DESC must be installed to load '
+            'a FourierCurrentPotentialField.'
+        )
+    mm, nn, phis_raw, phic_raw = ptolemy_identity_rev(scf.Phi_basis.modes[:,1], scf.Phi_basis.modes[:,2], scf.Phi_mn)
+    mpol = scf.M
+    ntor = scf.N
+    stellsym = scf.sym
+    nfp = scf.NFP
+    phic = phic_raw.flatten()
+    phis = phis_raw.flatten()[1:]
+    # When stellsym is enabled, Phi is a sin fourier series.
+    # Otherwise, it's a sin-cos series.
+    if stellsym:
+        phi = jnp.concatenate([phis])
+    else:
+        phi = jnp.concatenate([phis, phic])
+    if quadpoints_phi and quadpoints_theta:
+        if source_grid:
+            raise ValueError(
+                'The user should provide either '
+                '(quadpoints_phi, quadpoints_theta) or source_grid, '
+                'not both.'
+            )
+        winding_quadpoints_phi_temp = winding_quadpoints_phi
+        winding_quadpoints_theta_temp = winding_quadpoints_theta
+    else:
+        winding_quadpoints_phi_temp = jnp.linspace(0, 1, (source_grid.M*2+1) * nfp, endpoint=False)
+        winding_quadpoints_theta_temp = jnp.linspace(0, 1, source_grid.N*2+1, endpoint=False)
+    
+    winding_surface = SurfaceRZFourierJAX.from_desc(
+        scf_nescoil, 
+        quadpoints_phi = winding_quadpoints_phi_temp,
+        quadpoints_theta = winding_quadpoints_theta_temp,
+    )
+    return (qp, {'phi': phi})

quadcoil/io/focus.py

@@ -0,0 +1,65 @@
+import numpy as np
+from quadcoil import SurfaceRZFourierJAX
+def save_focus(surface, filename):
+    '''
+    Saves a Simsopt `SurfaceRZFourier`, a DESC `FourierRZToroidalSurface`,
+    or a QUADCOIL `SurfaceRZFourierJAX`
+    to a `.plasma` file used for FOCUS or FAMUS under 
+    `CASE_SURFACE=0`.
+
+    Parameters
+    ----------
+    surface
+        The surface tp save.
+    filename
+        The filename to save.
+    '''
+    try:
+        from simsopt.geo import SurfaceRZFourier
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError('Simsopt must be installed to output as FOCUS.')
+    if isinstance(surface, SurfaceRZFourierJAX):
+        surface = surface.to_simsopt()
+    elif isinstance(surface, SurfaceRZFourier):
+        pass
+    else:
+        # This order of the type circumvents the DESC import unless
+        # necessary.
+        from desc.geometry.surface import FourierRZToroidalSurface
+        if not isinstance(surface, FourierRZToroidalSurface):
+            raise TypeError(
+                'surface must be a Simsopt SurfaceRZFourier, '
+                'a DESC FourierRZToroidalSurface, '
+                'or a QUADCOIL `SurfaceRZFourierJAX`'
+            )
+        surface = SurfaceRZFourierJAX.from_desc(
+            surface,
+            quadpoints_phi=np.linspace(0, 1, 32, endpoint=False),
+            quadpoints_theta=np.linspace(0, 1, 32, endpoint=False)
+        )
+    
+    replace_err = lambda x: 0 if isinstance(x, ValueError) or not np.isfinite(x) else x
+    lines_to_write = [
+        '# N_modes N_fp  N_bn\n',
+        str(len(surface.get_dofs())) + '  ' + str(surface.nfp) + '  ' + '0\n', # No Bnorm modes  
+        '#------- plasma boundary------\n',
+        '#  n   m   Rbc   Rbs    Zbc   Zbs\n',
+    ]
+    for m in np.arange(0, surface.mpol+1):
+        for n in np.arange(-surface.ntor, surface.ntor+1):
+            Rbc = surface.get_rc(m, n)
+            Rbs = surface.get_rs(m, n)
+            Zbc = surface.get_zc(m, n)
+            Zbs = surface.get_zs(m, n)
+            lines_to_write+=[
+                str(n) + '  '
+                + str(m) + '  '
+                + str(replace_err(Rbc)) + '  '
+                + str(replace_err(Rbs)) + '  '
+                + str(replace_err(Zbc)) + '  '
+                + str(replace_err(Zbs)) + '\n'
+            ]
+    if not filename.endswith(".plasma"):
+        filename += ".plasma"
+    with open(filename, "w") as f:
+        f.writelines(lines_to_write)