Finite Build Coils

desc/compute/_curve.py

@@ -1028,3 +1028,72 @@ def _length_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
         # but also works if grid.endpoint is False
         data["length"] = jnp.sum(T * data["ds"])
     return data
+
+
+@register_compute_fun(
+    name="centroid_tangent",
+    label="\\mathbf{T}_{\\mathrm{Centroid}}",
+    units="~",
+    units_long="None",
+    description="Tangent unit vector to curve in centroid frame (see Singh et al. 2020, section 3.1)",
+    dim=3,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="s",
+    data=["x_s"],
+    parameterization="desc.geometry.core.Curve",
+)
+def _centroid_tangent(params, transforms, profiles, data, **kwargs):
+    # this is defined identically to the Frenet-Serret tangent
+    data["centroid_tangent"] = (
+        data["x_s"] / jnp.linalg.norm(data["x_s"], axis=-1)[:, None]
+    )
+    return data
+
+
+@register_compute_fun(
+    name="centroid_normal",
+    label="\\mathbf{N}_{\\mathrm{Centroid}}",
+    units="~",
+    units_long="None",
+    description="Normal unit vector to curve in centroid frame (see Singh et al. 2020, section 3.1)",
+    dim=3,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="s",
+    data=["x", "center", "centroid_tangent"],
+    parameterization="desc.geometry.core.Curve",
+)
+def _centroid_normal(params, transforms, profiles, data, **kwargs):
+    delta = data["x"] - data["center"]  # do I need to convert this to xyz?
+    delta_orth = (
+        delta - dot(delta, data["centroid_tangent"])[:, None] * data["centroid_tangent"]
+    )
+    delta_orth = delta_orth / jnp.linalg.norm(delta_orth, axis=-1)[:, None]
+    data["centroid_normal"] = delta_orth
+
+    return data
+
+
+@register_compute_fun(
+    name="centroid_binormal",
+    label="\\mathbf{B}_{\\mathrm{Centroid}}",
+    units="~",
+    units_long="None",
+    description="Binormal unit vector to curve in centroid frame (see Singh et al. 2020, section 3.1)",
+    dim=3,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="s",
+    data=["centroid_tangent", "centroid_normal"],
+    parameterization="desc.geometry.core.Curve",
+)
+def _centroid_binormal(params, transforms, profiles, data, **kwargs):
+    data["centroid_binormal"] = cross(
+        data["centroid_tangent"], data["centroid_normal"]
+    )  # TODO: figure out if rotation is necessary for these unit vecs
+
+    return data

desc/compute/_field.py

@@ -3516,3 +3516,323 @@ def _L_grad_B(params, transforms, profiles, data, **kwargs):
 def _K_vc(params, transforms, profiles, data, **kwargs):
     data["K_vc"] = cross(data["n_rho"], data["B"]) / mu_0
     return data
+
+
+@register_compute_fun(
+    name="p_frame",
+    label="\\mathbf{p}_{\\mathrm{Frame}}",
+    units="~",
+    units_long="None",
+    description="P unit vector for framed coil with arbitrary rotation fourier series",
+    dim=3,
+    params=["alpha_n"],
+    transforms={"alpha": [[0, 0, 0]]},
+    profiles=[],
+    coordinates="s",
+    data=["centroid_normal", "centroid_binormal"],
+    parameterization="desc.coils._FramedCoil",
+)
+def _p_frame(params, transforms, profiles, data, **kwargs):
+    alpha = transforms["alpha"].transform(
+        params["alpha_n"], dz=0
+    )  # TODO: check if this is right
+
+    p_frame = (
+        data["centroid_normal"] * jnp.cos(alpha)[:, jnp.newaxis]
+        + data["centroid_binormal"] * jnp.sin(alpha)[:, jnp.newaxis]
+    )
+
+    data["p_frame"] = p_frame
+    return data
+
+
+@register_compute_fun(
+    name="q_frame",
+    label="\\mathbf{q}_{\\mathrm{Frame}}",
+    units="~",
+    units_long="None",
+    description="Q unit vector for framed coil with arbitrary rotation fourier series",
+    dim=3,
+    params=["alpha_n"],
+    transforms={"alpha": [[0, 0, 0]]},
+    profiles=[],
+    coordinates="s",
+    data=["centroid_normal", "centroid_binormal"],
+    parameterization="desc.coils._FramedCoil",
+)
+def _q_frame(params, transforms, profiles, data, **kwargs):
+    alpha = transforms["alpha"].transform(params["alpha_n"], dz=0)
+
+    q_frame = (
+        -data["centroid_normal"] * jnp.sin(alpha)[:, jnp.newaxis]
+        + data["centroid_binormal"] * jnp.cos(alpha)[:, jnp.newaxis]
+    )
+
+    data["q_frame"] = q_frame
+    return data
+
+
+@register_compute_fun(
+    name="curv1_frame",
+    label="\\kappa_{1}}_{\\mathrm{Frame}",
+    units="~",
+    units_long="None",
+    description="Curvature component in the p-vector direction for a rectangular cross section coil",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="s",
+    data=["curvature", "frenet_normal", "p_frame"],
+    parameterization="desc.coils._FramedCoil",
+)
+def _curv1_frame(params, transforms, profiles, data, **kwargs):
+    data["curv1_frame"] = data["curvature"] * dot(
+        data["frenet_normal"], data["p_frame"]
+    )
+
+    return data
+
+
+@register_compute_fun(
+    name="curv2_frame",
+    label="\\kappa_{2}}_{\\mathrm{Frame}",
+    units="~",
+    units_long="None",
+    description="Curvature component in the q-vector direction for a rectangular cross section coil",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="s",
+    data=["curvature", "frenet_normal", "q_frame"],
+    parameterization="desc.coils._FramedCoil",
+)
+def _curv2_frame(params, transforms, profiles, data, **kwargs):
+    data["curv2_frame"] = data["curvature"] * dot(
+        data["frenet_normal"], data["q_frame"]
+    )
+
+    return data
+
+
+@register_compute_fun(
+    name="u_fb",
+    label="u_{FB}",
+    units="",
+    units_long="",
+    description="U coordinate for finite build expansion",
+    dim=1,
+    params=[],
+    transforms={"grid": []},
+    profiles=[],
+    coordinates="rt",  # just two degrees of freedom
+    data=[],
+    parameterization="desc.coils._FiniteBuildCoil",
+)
+def _u_fb(params, transforms, profiles, data, **kwargs):
+    grid = transforms["grid"]
+    # TODO: check that grid is LinearGrid?
+    data["u_fb"] = (
+        (grid.nodes[:, 0] - 0.5) / (0.5) * 0.999
+    )  # rescale to [-0.999,0.999], as the finite build term is singular at the edge
+    return data
+
+
+@register_compute_fun(
+    name="v_fb",
+    label="v_{FB}",
+    units="",
+    units_long="",
+    description="V coordinate for finite build expansion",
+    dim=1,
+    params=[],
+    transforms={"grid": []},
+    profiles=[],
+    coordinates="rt",
+    data=[],
+    parameterization="desc.coils._FiniteBuildCoil", 
+)
+def _v_fb(params, transforms, profiles, data, **kwargs):
+    grid = transforms["grid"]
+    # TODO: check that grid is LinearGrid?
+    data["v_fb"] = (
+        (grid.nodes[:, 1] - jnp.pi) / (jnp.pi) * 0.999
+    )  # rescale to [-0.999,0.999]
+    return data
+
+
+@register_compute_fun(
+    name="B_0_fb",
+    label="B_{0,FB}",
+    units="T",
+    units_long="Tesla",
+    description="B_0 term in finite build expansion",
+    dim=3,
+    params=["p_frame", "q_frame", "current", "cross_section_dims"],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["u_fb", "v_fb"],
+    parameterization="desc.coils._FiniteBuildCoil",
+)
+def _B_0_fb(params, transforms, profiles, data, **kwargs):
+    # fed in externally
+    p_frame = params["p_frame"]
+    q_frame = params["q_frame"]
+
+    # scalars
+    current = params["current"]
+    a = params["cross_section_dims"][0]
+    b = params["cross_section_dims"][1]
+
+    # u and v are computed
+    u = data["u_fb"]
+    v = data["v_fb"]
+
+    # add dimension to u and v to accomodate vectorized operations
+    u = u[:, jnp.newaxis]
+    v = v[:, jnp.newaxis]
+
+    def G(x, y):
+        return y * jnp.arctan(x / y) + x / 2 * jnp.log(1 + (y / x) ** 2)
+
+    data["B_0_fb"] = mu_0 * current/(4*jnp.pi*a*b) * ( (q_frame * G(b*(v+1), a*(u+1)) - p_frame * G(a*(u+1), b*(v+1))) + 
+                                (-1) * (q_frame * G(b*(v+1), a*(u-1)) - p_frame * G(a*(u-1), b*(v+1))) + 
+                                (-1) * (q_frame * G(b*(v-1), a*(u+1)) - p_frame * G(a*(u+1), b*(v-1))) + 
+                                (q_frame * G(b*(v-1), a*(u-1)) - p_frame * G(a*(u-1), b*(v-1))))
+
+    return data
+
+
+@register_compute_fun(
+    name="B_kappa_fb",
+    label="B_{\\kappa,FB}",
+    units="T",
+    units_long="Tesla",
+    description="B_kappa term in finite build expansion",
+    dim=3,
+    params=[
+        "p_frame",
+        "q_frame",
+        "current",
+        "cross_section_dims",
+        "curv1_frame",
+        "curv2_frame",
+    ],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["u_fb", "v_fb"],
+    parameterization="desc.coils._FiniteBuildCoil", 
+)
+def _B_kappa_fb(params, transforms, profiles, data, **kwargs):
+    # fed in externally
+    p_frame = params["p_frame"]
+    q_frame = params["q_frame"]
+    curv1_frame = params["curv1_frame"]
+    curv2_frame = params["curv2_frame"]
+
+    # scalars
+    current = params["current"]
+    a = params["cross_section_dims"][0]
+    b = params["cross_section_dims"][1]
+
+    # u and v are computed
+    u = data["u_fb"]
+    v = data["v_fb"]
+
+    # add dimension to scalars to accomodate vectorized operations
+    u = u[:, jnp.newaxis]
+    v = v[:, jnp.newaxis]
+    curv1_frame = curv1_frame[:, jnp.newaxis]
+    curv2_frame = curv2_frame[:, jnp.newaxis]
+
+    def K(U,V):
+        return ((-2*U*V * (curv1_frame*q_frame-curv2_frame*p_frame) * jnp.log(a/b*U**2 + b/a*V**2) + 
+                (curv2_frame*q_frame-curv1_frame*p_frame) * (a/b*U**2 + b/a*V**2) * jnp.log(a/b*U**2 + b/a*V**2) + 
+                4*a/b*curv2_frame*p_frame*U**2 * jnp.arctan(b*V/(a*U)) -
+                4*b/a*curv1_frame*q_frame*V**2 * jnp.arctan(a*U/(b*V))) )
+
+    data["B_kappa_fb"] = mu_0 * current/(64*jnp.pi) * ( K(u+1, v+1) + K(u-1, v-1) +
+                                  (-1) * K(u-1, v+1) + (-1) * K(u+1, v-1))
+
+    return data
+
+
+@register_compute_fun(
+    name="B_b_fb",
+    label="B_{b,FB}",
+    units="T",
+    units_long="Tesla",
+    description="B_b term in finite build expansion",
+    dim=3,
+    params=["current", "cross_section_dims"],
+    transforms={},
+    profiles=[],
+    coordinates="r",
+    data=["curvature", "frenet_binormal"],
+    parameterization="desc.coils._FiniteBuildCoil",  
+)
+def _B_b_fb(params, transforms, profiles, data, **kwargs):
+    # scalars, external
+    current = params["current"]
+    a = params["cross_section_dims"][0]
+    b = params["cross_section_dims"][1]
+
+    # u and v are computed
+    curvature = data["curvature"][:, jnp.newaxis]
+    binormal = data["frenet_binormal"]
+
+    k = -(a**4 - 6 * a**2 * b**2 + b**4) / (6 * a**2 * b**2) * jnp.log(a / b + b / a)
+    +b * b / (6 * a * a) * jnp.log(b / a)
+    +a * a / (6 * b * b) * jnp.log(a / b)
+    +(4.0 * b) / (3 * a) * jnp.arctan(a / b)
+    +(4.0 * a) / (3 * b) * jnp.arctan(b / a)
+
+    delta = jnp.exp(-(25.0 / 6) + k)
+
+    data["B_b_fb"] = mu_0 * current/(8*jnp.pi) * curvature * binormal * (4 + 2*jnp.log(2) + jnp.log(delta))
+
+    return data
+
+@register_compute_fun(
+    name="x_fb",
+    label="x_{FB}",
+    units="m",
+    units_long="meters",
+    description="Position vector in lab frame for finite build coil cross section",
+    dim=3,
+    params=[
+        "p_frame",
+        "q_frame",
+        "cross_section_dims",
+        "x_centerline"
+    ],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["u_fb", "v_fb"],
+    parameterization="desc.coils._FiniteBuildCoil", 
+)
+def _x_fb(params, transforms, profiles, data, **kwargs):
+    # fed in externally
+    p_frame = params["p_frame"]
+    q_frame = params["q_frame"]
+    x_centerline = params["x_centerline"]
+
+    # scalars
+    a = params["cross_section_dims"][0]
+    b = params["cross_section_dims"][1]
+
+    # u and v are computed
+    u = data["u_fb"]
+    v = data["v_fb"]
+
+    # add dimension to scalars to accomodate vectorized operations
+    u = u[:, jnp.newaxis]
+    v = v[:, jnp.newaxis]
+
+    data["x_fb"] = x_centerline + a * u * p_frame + b * v * q_frame
+
+    return data