diff --git a/essos/coil_perturbation.py b/essos/coil_perturbation.py index 3be91e33..3d8b4232 100644 --- a/essos/coil_perturbation.py +++ b/essos/coil_perturbation.py @@ -3,7 +3,7 @@ import jax.numpy as jnp from jax import jit, vmap from jaxtyping import Array, Float # https://github.com/google/jaxtyping -from essos.coils import Curves,apply_symmetries_to_gammas +from essos.coils import Curves, Coils, DiscretizedCoils, fit_dofs_from_coils from functools import partial @@ -205,70 +205,94 @@ def get_sample(self, deriv): -def perturb_curves_systematic(curves: Curves,sampler:GaussianSampler, key=None): - """ - Apply a systematic perturbation to all the coils. - This means taht an independent perturbation is applied to the each unique coil - Then, the required symmetries are applied to the perturbed unique set of coils - - Args: - curves: curves to be perturbed. - sampler: the gaussian sampler used to get the perturbations - key: the seed which will be splited to geenerate random - but reproducible pertubations - - Returns: - The curves given as an input are modified and thus no return is done - """ - new_seeds=jax.random.split(key, num=curves.n_base_curves) - if sampler.n_derivs == 0: - perturbation = jax.vmap(sampler.draw_sample, in_axes=(0))(new_seeds) - gamma_perturbations = apply_symmetries_to_gammas(perturbation[:,0,:,:], curves.nfp, curves.stellsym) - curves.gamma=curves.gamma + gamma_perturbations - elif sampler.n_derivs == 1: - perturbation = jax.vmap(sampler.draw_sample, in_axes=(0))(new_seeds) - gamma_perturbations = apply_symmetries_to_gammas(perturbation[:,0,:,:], curves.nfp, curves.stellsym) - gamma_perturbations_dash = apply_symmetries_to_gammas(perturbation[:,1,:,:], curves.nfp, curves.stellsym) - curves.gamma=curves.gamma + gamma_perturbations - curves.gamma_dash=curves.gamma_dash + gamma_perturbations_dash - elif sampler.n_derivs == 2: - perturbation = jax.vmap(sampler.draw_sample, in_axes=(0))(new_seeds) - gamma_perturbations = apply_symmetries_to_gammas(perturbation[:,0,:,:], curves.nfp, curves.stellsym) - gamma_perturbations_dash = apply_symmetries_to_gammas(perturbation[:,1,:,:], curves.nfp, curves.stellsym) - gamma_perturbations_dashdash = apply_symmetries_to_gammas(perturbation[:,2,:,:], curves.nfp, curves.stellsym) - curves.gamma=curves.gamma + gamma_perturbations - curves.gamma_dash=curves.gamma_dash + gamma_perturbations_dash - curves.gamma_dashdash=curves.gamma_dashdash + gamma_perturbations_dashdash - #return curves - - -def perturb_curves_statistic(curves: Curves,sampler:GaussianSampler, key=None): +#Util functions for perturbing curves and coils. +def _draw_curve_perturbation(sampler: GaussianSampler, key, n_curves: int): + new_seeds = jax.random.split(key, num=n_curves) + perturbation = jax.vmap(sampler.draw_sample, in_axes=(0))(new_seeds) + return perturbation[:, 0, :, :] + + +def _make_curves_like(curves: Curves, dofs_new, nfp: int, stellsym: bool): + return Curves( + dofs_new, + curves.n_segments, + nfp=nfp, + stellsym=stellsym, + scaling_type=curves.scaling_type, + scaling_factor=curves.scaling_factor, + scale_fixed=curves.scale_fixed, + ) + + + +#Perturb coisl fucntion +def perturb_curves(curves, sampler:GaussianSampler, key=None, perturbation_type="systematic"): """ - Apply a statistic perturbation to all the coils. - This means taht an independent perturbation is applied every coil - including repeated coils - + Apply a perturbation to curves or coils and return the perturbed object. + Args: - curves: curves to be perturbed. + curves: Curves, Coils, or DiscretizedCoils to be perturbed. sampler: the gaussian sampler used to get the perturbations - key: the seed which will be splited to geenerate random - but reproducible pertubations - + key: the seed which will be split to generate random but reproducible perturbations + perturbation_type: "systematic" to perturb only unique/base coils and preserve symmetry, + or "statistical"/"statistic" to perturb every expanded coil independently. + Returns: - The curves given as an input are modified and thus no return is done + A new perturbed object of the same family as the input. """ - new_seeds=jax.random.split(key, num=curves.gamma.shape[0]) - if sampler.n_derivs == 0: - perturbation = jax.vmap(sampler.draw_sample, in_axes=(0))(new_seeds) - curves.gamma=curves.gamma + perturbation[:,0,:,:] - elif sampler.n_derivs == 1: - perturbation = jax.vmap(sampler.draw_sample, in_axes=(0))(new_seeds) - curves.gamma=curves.gamma + perturbation[:,0,:,:] - curves.gamma_dash=curves.gamma_dash + perturbation[:,1,:,:] - elif sampler.n_derivs == 2: - perturbation = jax.vmap(sampler.draw_sample, in_axes=(0))(new_seeds) - curves.gamma=curves.gamma + perturbation[:,0,:,:] - curves.gamma_dash=curves.gamma_dash + perturbation[:,1,:,:] - curves.gamma_dashdash=curves.gamma_dashdash + perturbation[:,2,:,:] - #return curves - + if perturbation_type == "systematic": + if isinstance(curves, DiscretizedCoils): + perturbation = _draw_curve_perturbation(sampler, key, curves.n_base_curves) + return DiscretizedCoils( + curves.dofs_gamma + perturbation, + currents=curves.dofs_currents_raw, + nfp=curves.nfp, + stellsym=curves.stellsym, + ) + + if isinstance(curves, Coils): + perturbation = _draw_curve_perturbation(sampler, key, curves.curves.n_base_curves) + base_curves = _make_curves_like(curves.curves, curves.dofs_curves, nfp=1, stellsym=False) + perturbed_base_gamma = base_curves.gamma + perturbation + dofs_new, _ = fit_dofs_from_coils(perturbed_base_gamma, curves.order, curves.n_segments, assume_uniform=True) + new_curves = _make_curves_like(curves.curves, dofs_new, nfp=curves.nfp, stellsym=curves.stellsym) + return Coils(curves=new_curves, currents=curves.dofs_currents_raw) + + if isinstance(curves, Curves): + perturbation = _draw_curve_perturbation(sampler, key, curves.n_base_curves) + base_curves = _make_curves_like(curves, curves.dofs, nfp=1, stellsym=False) + perturbed_base_gamma = base_curves.gamma + perturbation + dofs_new, _ = fit_dofs_from_coils(perturbed_base_gamma, curves.order, curves.n_segments, assume_uniform=True) + return _make_curves_like(curves, dofs_new, nfp=curves.nfp, stellsym=curves.stellsym) + + elif perturbation_type in {"statistical"}: + perturbation = _draw_curve_perturbation(sampler, key, curves.gamma.shape[0]) + gamma_perturbed = curves.gamma + perturbation + + if isinstance(curves, DiscretizedCoils): + return DiscretizedCoils(gamma_perturbed, currents=curves.currents, nfp=1, stellsym=False) + + if isinstance(curves, Coils): + dofs_new, _ = fit_dofs_from_coils(gamma_perturbed, curves.order, curves.n_segments, assume_uniform=True) + new_curves = _make_curves_like(curves.curves, dofs_new, nfp=1, stellsym=False) + return Coils(curves=new_curves, currents=curves.currents) + + if isinstance(curves, Curves): + dofs_new, _ = fit_dofs_from_coils(gamma_perturbed, curves.order, curves.n_segments, assume_uniform=True) + return _make_curves_like(curves, dofs_new, nfp=1, stellsym=False) + + else: + raise ValueError( + f"Unsupported perturbation_type {perturbation_type}. " + "Expected 'systematic' or 'statistical'." + ) + + raise TypeError(f"Unsupported type {type(curves)}. Expected Curves, Coils, or DiscretizedCoils.") + + +def perturb_curves_systematic(curves, sampler:GaussianSampler, key=None): + return perturb_curves(curves, sampler, key=key, perturbation_type="systematic") + + +def perturb_curves_statistic(curves, sampler:GaussianSampler, key=None): + return perturb_curves(curves, sampler, key=key, perturbation_type="statistical") diff --git a/essos/coils.py b/essos/coils.py index 76955ea1..c4c4c3cd 100644 --- a/essos/coils.py +++ b/essos/coils.py @@ -25,7 +25,30 @@ def __init__(self, dofs: jnp.ndarray, n_segments: int = 100, nfp: int = 1, - stellsym: bool = True): + stellsym: bool = True, + scaling_type: int = 2, + scaling_factor: float = 0.0, + scale_fixed: float = 1.0): + """Initialize Curves. + + Args: + dofs: Fourier coefficients with shape ``(n_curves, 3, 2*order+1)``. + n_segments: number of quadrature points used to discretize each curve. + nfp: number of field periods. + stellsym: whether stellarator symmetry is used. + scaling_type: norm used in the mode scaling. Accepted values are + ``'L1'`` or ``1``, ``'L2'`` or ``2``, and ``'Linfty'`` or ``-1``. + scaling_factor: exponential weight used in the scaling + ``exp(scaling_factor * ||mode_orders||)``. + scale_fixed: fixed multiplier applied to all modes. + + Note: + The optimized dofs are stored as ``_dofs * scaling``, while the + internal physical coefficients are kept in ``_dofs``. + The scaling interface matches the surface scaling options, but + here the norm is applied to a 1D mode-order measure, so ``L1``, + ``L2``, and ``Linfty`` currently give the same numerical scaling. + """ if hasattr(dofs, 'shape'): assert len(dofs.shape) == 3, "dofs must be a 3D array with shape (n_curves, 3, 2*order+1)" assert dofs.shape[1] == 3, "dofs must have shape (n_curves, 3, 2*order+1)" @@ -41,6 +64,11 @@ def __init__(self, self._nfp = nfp self._stellsym = stellsym + self._scaling_type = self._normalize_scaling_type(scaling_type) + self._scaling_factor = scaling_factor + self._scale_fixed = scale_fixed + self._scaling = None + self.quadpoints = jnp.linspace(0, 1, self._n_segments, endpoint=False) self._curves = None self._gamma = None @@ -48,6 +76,29 @@ def __init__(self, self._gamma_dashdash = None self._length = None self._curvature = None + + @staticmethod + def _normalize_scaling_type(scaling_type): + """Map public scaling_type inputs to norm orders used internally.""" + if scaling_type == "L1" or scaling_type == 1: + return 1 + if scaling_type == "L2" or scaling_type == 2: + return 2 + if scaling_type == "Linfty" or scaling_type == -1 or scaling_type == jnp.inf: + return jnp.inf + raise ValueError( + f"Unknown scaling_type: {scaling_type}. " + "Expected 'L1', 1, 'L2', 2, 'Linfty', -1, or jnp.inf." + ) + + @staticmethod + def _compute_mode_scaling(order, scaling_type, scaling_factor, scale_fixed): + mode_orders = jnp.concatenate([ + jnp.array([0.0]), + jnp.repeat(jnp.arange(1, order + 1, dtype=float), 2) + ]) + mode_norm = jnp.linalg.norm(jnp.vstack([mode_orders]), ord=scaling_type, axis=0) + return jnp.exp(scaling_factor * mode_norm) * scale_fixed # reset_cache method def reset_cache(self): @@ -61,12 +112,13 @@ def reset_cache(self): # dofs property and setter @property def dofs(self): - return jnp.array(self._dofs) + # Apply scaling to each coordinate (X, Y, Z) independently + return self._dofs * self.scaling[None, None, :] @dofs.setter def dofs(self, new_dofs): self.reset_cache() - self._dofs = new_dofs + self._dofs = new_dofs / self.scaling[None, None, :] # n_segments property and setter @property @@ -99,15 +151,59 @@ def stellsym(self, new_stellsym): self.reset_cache() self._stellsym = new_stellsym + # scaling_type property and setter + @property + def scaling_type(self): + return self._scaling_type + + @scaling_type.setter + def scaling_type(self, new_type): + self._scaling_type = self._normalize_scaling_type(new_type) + self._scaling = None + + # scaling_factor property and setter + @property + def scaling_factor(self): + return self._scaling_factor + + @scaling_factor.setter + def scaling_factor(self, new_factor): + self._scaling_factor = new_factor + self._scaling = None + + # scale_fixed property and setter + @property + def scale_fixed(self): + return self._scale_fixed + + @scale_fixed.setter + def scale_fixed(self, new_scale): + self._scale_fixed = new_scale + self._scaling = None + + # scaling property + @property + def scaling(self): + """Mode-by-mode scaling ``scale_fixed * exp(scaling_factor * ||mode_orders||)``.""" + if self._scaling is None: + self._scaling = self._compute_mode_scaling( + self.order, self.scaling_type, self.scaling_factor, self.scale_fixed + ) + return self._scaling + # order property and setter @property def order(self): - return self.dofs.shape[2]//2 + return self._dofs.shape[2]//2 @order.setter def order(self, new_order): self.reset_cache() - self._dofs = jnp.pad(self.dofs, ((0,0), (0,0), (0, max(0, 2*(new_order-self.order)))))[:, :, :2*(new_order)+1] + # Get unscaled dofs, resize, then store unscaled + old_scaling = self.scaling + unscaled_dofs = self._dofs + self._dofs = jnp.pad(unscaled_dofs, ((0,0), (0,0), (0, max(0, 2*(new_order-self.order)))))[:, :, :2*(new_order)+1] + self._scaling = None # Force recalculation for new order # n_base_curves property @property @@ -118,7 +214,8 @@ def n_base_curves(self): @property def curves(self): if self._curves is None: - self._curves = apply_symmetries_to_curves(self.dofs, self.nfp, self.stellsym) + # Use unscaled dofs for physical curve representation + self._curves = apply_symmetries_to_curves(self._dofs, self.nfp, self.stellsym) return self._curves # _compute_gamma method @@ -135,15 +232,9 @@ def create_data(order: int) -> jnp.ndarray: # gamma property @property def gamma(self): - # Allow downstream code (e.g. coil_perturbation) to override the - # computed gamma; otherwise compute from Fourier coefficients. - if self._gamma is not None: - return self._gamma - return self._compute_gamma() - - @gamma.setter - def gamma(self, value): - self._gamma = value + if self._gamma is None: + self._gamma = self._compute_gamma() + return self._gamma # _compute_gamma_dash method @jit @@ -157,13 +248,9 @@ def create_data(order: int) -> jnp.ndarray: # gamma_dash property @property def gamma_dash(self): - if self._gamma_dash is not None: - return self._gamma_dash - return self._compute_gamma_dash() - - @gamma_dash.setter - def gamma_dash(self, value): - self._gamma_dash = value + if self._gamma_dash is None: + self._gamma_dash = self._compute_gamma_dash() + return self._gamma_dash # _compute_gamma_dashdash method @jit @@ -177,13 +264,9 @@ def create_data(order: int) -> jnp.ndarray: # gamma_dashdash property @property def gamma_dashdash(self): - if self._gamma_dashdash is not None: - return self._gamma_dashdash - return self._compute_gamma_dashdash() - - @gamma_dashdash.setter - def gamma_dashdash(self, value): - self._gamma_dashdash = value + if self._gamma_dashdash is None: + self._gamma_dashdash = self._compute_gamma_dashdash() + return self._gamma_dashdash # length property @property @@ -343,10 +426,21 @@ def wrap(data): polyLinesToVTK(str(filename), np.array(x), np.array(y), np.array(z), pointsPerLine=np.array(ppl), pointData=pointData) @classmethod - def from_simsopt(cls, simsopt_curves, nfp=1, stellsym=True): + def from_simsopt(cls, simsopt_curves, nfp=1, stellsym=True, scaling_type=2, scaling_factor=0.0, scale_fixed=1.0): """ Create a Curves object from a list of simsopt curves. This assumes curves have all nfp and stellsym symmetries. + + Args: + scaling_type: accepted values are ``'L1'`` or ``1``, ``'L2'`` or ``2``, + and ``'Linfty'`` or ``-1``. + scaling_factor: exponential weight used in the mode scaling. + scale_fixed: fixed multiplier applied to all modes. + + Note: + The norm choice is kept consistent with surfaces, but for the + current 1D mode-order scaling it does not change the numerical + scaling. """ if isinstance(simsopt_curves, str): from simsopt import load @@ -358,23 +452,47 @@ def from_simsopt(cls, simsopt_curves, nfp=1, stellsym=True): [curve.x for curve in simsopt_curves] ), (len(simsopt_curves), 3, 2*simsopt_curves[0].order+1)) n_segments = len(simsopt_curves[0].quadpoints) - return cls(dofs, n_segments, nfp, stellsym) + return cls(dofs, n_segments, nfp, stellsym, scaling_type, scaling_factor, scale_fixed) def _tree_flatten(self): - children = (self._dofs,) # arrays / dynamic values + children = (self.dofs,) # arrays / dynamic values aux_data = {"n_segments": self._n_segments, "nfp": self._nfp, - "stellsym": self._stellsym} # static values + "stellsym": self._stellsym, + "scaling_type": self._scaling_type, + "scaling_factor": self._scaling_factor, + "scale_fixed": self._scale_fixed} # static values return (children, aux_data) @classmethod def _tree_unflatten(cls, aux_data, children): - return cls(*children, **aux_data) + dofs, = children + order = dofs.shape[2] // 2 + scaling_type = cls._normalize_scaling_type(aux_data["scaling_type"]) + scaling = cls._compute_mode_scaling( + order, scaling_type, aux_data["scaling_factor"], aux_data["scale_fixed"] + ) + return cls(dofs / scaling[None, None, :], **aux_data) tree_util.register_pytree_node(Curves, Curves._tree_flatten, Curves._tree_unflatten) + +def _initialize_currents_scale(currents, currents_scale): + """Return a fixed current scale for normalized current dofs.""" + if currents_scale is None: + return jnp.mean(jnp.abs(currents)) + return currents_scale + + +def _initialize_scale_fixed(gamma, scale_fixed): + """Return a fixed geometry scale for normalized gamma dofs.""" + if scale_fixed is None: + return jnp.maximum(jnp.max(jnp.abs(gamma)), 1.0) + return scale_fixed + + # TODO: change currents logic: save dofs_currents as dynamic -> alter main class Coils: """ Class to store the coils @@ -389,21 +507,28 @@ class Coils: dofs (jnp.ndarray - shape (n_base_curves * 3 * (2 * order + 1) + n_base_curves,)): Degrees of freedom of the coils (curves and normalized currents) """ - def __init__(self, curves: Curves, currents: jnp.ndarray): + def __init__(self, curves: Curves, currents: jnp.ndarray, currents_scale=None): + """Initialize coils. + + Args: + curves: base curve geometry. + currents: raw physical base currents. + currents_scale: fixed normalization used for ``dofs_currents``. + If ``None``, it is computed once from ``currents`` and then kept fixed. + """ # if hasattr(curves, 'n_base_curves') and hasattr(currents, 'size'): # assert curves.n_base_curves == currents.size, "Number of base curves and number of currents must be the same" self.curves = curves self._dofs_currents_raw = currents # Non-normalized base currents - self._currents_scale = None + self._currents_scale = _initialize_currents_scale(currents, currents_scale) self._dofs_currents = None self._currents = None # reset_cache method def reset_cache(self): self._dofs_currents = None - self._currents_scale = None self._currents = None # dofs_curves property and setter @@ -428,8 +553,6 @@ def dofs_currents_raw(self, new_dofs_currents_raw): # currents_scale property and setter @property def currents_scale(self): - if self._currents_scale is None: - self._currents_scale = jnp.mean(jnp.abs(self.dofs_currents_raw)) return self._currents_scale @currents_scale.setter @@ -535,11 +658,10 @@ def n_segments(self, new_n_segments): # copy method def copy(self): - coils = Coils(self.curves.copy(), self.dofs_currents_raw.copy()) + coils = Coils(self.curves.copy(), self.dofs_currents_raw.copy(), currents_scale=self.currents_scale) # Initialize caches coils._dofs_currents = self.dofs_currents - coils._currents_scale = self.currents_scale coils._currents = self._currents return coils @@ -601,6 +723,12 @@ def save_coils(self, filename: str, text=""): file.write(f"{self.nfp} {self.stellsym} {self.order}\n") file.write(f"Degrees of freedom\n") file.write(f"{repr(self.dofs.tolist())}\n") + file.write(f"Curves scaling type\n") + file.write(f"{self.curves.scaling_type}\n") + file.write(f"Curves scaling factor\n") + file.write(f"{self.curves.scaling_factor}\n") + file.write(f"Curves fixed scaling\n") + file.write(f"{self.curves.scale_fixed}\n") file.write(f"Currents degrees of freedom\n") file.write(f"{repr(self._dofs_currents.tolist())}\n") file.write(f"Currents scaling factor\n") @@ -620,17 +748,30 @@ def to_simsopt(self): return coils_via_symmetries(cuves_simsopt, currents_simsopt, self.nfp, self.stellsym) def to_json(self, filename: str): + """Save coils to JSON with proper scaling metadata. + + Saves raw unscaled DOFs (_dofs) along with all scaling parameters + to ensure perfect reconstruction on load. + """ data = { "nfp": self.nfp, "stellsym": self.stellsym, "order": self.order, "n_segments": self.n_segments, - "dofs_curves": self.dofs_curves.tolist(), - "dofs_currents": self.dofs_currents.tolist(), + # Save RAW unscaled curve DOFs + "dofs_curves_raw": jnp.asarray(self.curves._dofs).tolist(), + # Save curve scaling metadata + "scaling_type": self.curves.scaling_type, + "scaling_factor": float(self.curves.scaling_factor), + "scale_fixed": float(self.curves.scale_fixed), + # Save RAW unscaled currents + "dofs_currents_raw": jnp.asarray(self._dofs_currents_raw).tolist(), + # Save current scale if computed (optional for backward compat) + "currents_scale": float(self.currents_scale) if self._currents_scale is not None else None, } import json with open(filename, 'w') as file: - json.dump(data, file) + json.dump(data, file, indent=2) def plot(self, *args, **kwargs): self.curves.plot(*args, **kwargs) @@ -639,34 +780,82 @@ def to_vtk(self, *args, **kwargs): self.curves.to_vtk(*args, **kwargs) @classmethod - def from_simsopt(cls, simsopt_coils, nfp=1, stellsym=True): - """ This assumes coils have all nfp and stellsym symmetries""" + def from_simsopt(cls, simsopt_coils, nfp=1, stellsym=True, scaling_type=2, scaling_factor=0.0, scale_fixed=1.0): + """Create coils from simsopt coils. + + This assumes coils have all nfp and stellsym symmetries. + + Args: + scaling_type: accepted values are ``'L1'`` or ``1``, ``'L2'`` or ``2``, + and ``'Linfty'`` or ``-1``. + scaling_factor: exponential weight used in the mode scaling. + scale_fixed: fixed multiplier applied to all curve modes. + """ if isinstance(simsopt_coils, str): from simsopt import load bs = load(simsopt_coils) simsopt_coils = bs.coils curves = [c.curve for c in simsopt_coils] currents = jnp.array([c.current.get_value() for c in simsopt_coils[0:int(len(simsopt_coils)/nfp/(1+stellsym))]]) - return cls(Curves.from_simsopt(curves, nfp, stellsym), currents) + return cls(Curves.from_simsopt(curves, nfp, stellsym, scaling_type, scaling_factor, scale_fixed), currents) @classmethod def from_json(cls, filename: str): - """ Creates a Coils object from a json file""" + """Load coils from JSON with proper scaling metadata. + + Supports both new format (with raw DOFs and scaling) and legacy format + (with scaled DOFs) for backward compatibility. The scaling metadata + includes ``scaling_type``, ``scaling_factor``, and ``scale_fixed``. + """ import json with open(filename, "r") as file: data = json.load(file) - curves = Curves(jnp.array(data["dofs_curves"]), data["n_segments"], data["nfp"], data["stellsym"]) - currents = jnp.array(data["dofs_currents"]) - return cls(curves, currents) + + # Extract scaling metadata (with defaults for legacy files) + scaling_type = data.get("scaling_type", 2) + scaling_factor = data.get("scaling_factor", 0.0) + scale_fixed = data.get("scale_fixed", 1.0) + + # Check if using NEW format (raw DOFs) or LEGACY format (scaled DOFs) + if "dofs_curves_raw" in data: + # NEW FORMAT: Raw unscaled DOFs with full metadata + curves = Curves( + jnp.array(data["dofs_curves_raw"]), # Raw _dofs + data["n_segments"], + data["nfp"], + data["stellsym"], + scaling_type, + scaling_factor, + scale_fixed + ) + currents_raw = jnp.array(data["dofs_currents_raw"]) + else: + # LEGACY FORMAT: Assume "dofs_curves" are raw DOFs (old behavior) + # This maintains backward compatibility with old JSON files + curves = Curves( + jnp.array(data["dofs_curves"]), # Treat as raw for legacy + data["n_segments"], + data["nfp"], + data["stellsym"], + scaling_type, + scaling_factor, + scale_fixed + ) + # Legacy files may have scaled or raw currents - treat as raw + currents_raw = jnp.array(data["dofs_currents"]) + + # Create Coils object with raw currents + return cls(curves, currents_raw, currents_scale=data.get("currents_scale", None)) def _tree_flatten(self): - children = (self.curves, self._dofs_currents_raw) # arrays / dynamic values - aux_data = {} # static values + children = (self.curves, self.dofs_currents) # arrays / dynamic values + aux_data = {"currents_scale": self.currents_scale} # static values return (children, aux_data) @classmethod def _tree_unflatten(cls, aux_data, children): - return cls(*children, **aux_data) + curves, dofs_currents = children + return cls(curves, dofs_currents * aux_data["currents_scale"], currents_scale=aux_data["currents_scale"]) tree_util.register_pytree_node(Coils, Coils._tree_flatten, @@ -679,9 +868,23 @@ def CreateEquallySpacedCurves(n_curves: int, r: float, n_segments: int = 100, nfp: int = 1, - stellsym: bool = False) -> Curves: + stellsym: bool = False, + scaling_type: int = 2, + scaling_factor: float = 0, + scale_fixed: float = 1.0) -> Curves: """ Creates n_curves equally spaced on a torus of major radius R and minor radius r using Fourier - representation up to the specified order.""" + representation up to the specified order. + + Args: + scaling_type: accepted values are ``'L1'`` or ``1``, ``'L2'`` or ``2``, + and ``'Linfty'`` or ``-1``. + scaling_factor: exponential weight used in the mode scaling. + scale_fixed: fixed multiplier applied to all modes. + + Note: + The norm choice is kept consistent with surfaces, but for the current + 1D mode-order scaling it does not change the numerical scaling. + """ angles = (jnp.arange(n_curves) + 0.5) * (2 * jnp.pi) / ((1 + int(stellsym)) * nfp * n_curves) curves = jnp.zeros((n_curves, 3, 1 + 2 * order)) @@ -690,7 +893,10 @@ def CreateEquallySpacedCurves(n_curves: int, curves = curves.at[:, 1, 0].set(jnp.sin(angles) * R) # y[0] curves = curves.at[:, 1, 2].set(jnp.sin(angles) * r) # y[2] curves = curves.at[:, 2, 1].set(-r) # z[1] (constant for all) - return Curves(curves, n_segments=n_segments, nfp=nfp, stellsym=stellsym) + return Curves(curves, n_segments=n_segments, nfp=nfp, stellsym=stellsym, scaling_type=scaling_type, scaling_factor=scaling_factor, scale_fixed=scale_fixed) + + + @partial(jit, static_argnames=["flip"]) def RotatedCurve(curve, phi, flip): @@ -834,4 +1040,585 @@ def fit_dofs_from_coils( gamma_uni = _resample_closed_curve_uniform_batch(coils_gamma, n_segments) # arclength (vmapped) dofs = _fit_real_fourier_batch(gamma_uni, order) # rFFT-based fit - return dofs, gamma_uni \ No newline at end of file + return dofs, gamma_uni + +class DiscretizedCoils: + """ Class to store coils from gamma (discretized curve coordinates) instead of Fourier coefficients + + This class is compatible with the Coils class but stores dofs as the actual gamma values + rather than Fourier expansion coefficients. Derivatives are computed numerically. + + Attributes: + dofs_gamma (jnp.ndarray - shape (n_base_curves, n_segments, 3)): Base discretized curves (dofs) + gamma (jnp.ndarray - shape (n_curves, n_segments, 3)): Discretized curves after symmetry expansion + currents (jnp.ndarray - shape (n_curves,)): Currents after symmetry expansion + n_segments (int): Number of segments in the discretization + nfp (int): Number of field periods + stellsym (bool): Stellarator symmetry + dofs_currents_raw (jnp.ndarray - shape (n_base_curves,)): Non-normalized base currents + currents_scale (float): Normalization factor for the currents + dofs_currents (jnp.ndarray - shape (n_base_curves,)): Normalized base currents + """ + def __init__(self, gamma: jnp.ndarray, currents: jnp.ndarray, nfp: int = 1, stellsym: bool = False, currents_scale=None, scale_fixed=None): + """ + Initialize DiscretizedCoils with discretized curve coordinates and currents, applying symmetries if possible. + Args: + gamma: shape (n_base_curves, n_segments, 3) - base discretized curve coordinates + currents: shape (n_base_curves,) - base currents for each unique curve + nfp: Number of field periods (default: 1) + stellsym: Stellarator symmetry (default: False) + currents_scale: fixed normalization used for ``dofs_currents``. + If ``None``, it is computed once from ``currents`` and then kept fixed. + scale_fixed: fixed normalization used for ``dofs_gamma``. + If ``None``, it is computed once from ``max(abs(gamma))`` and then kept fixed. + """ + gamma = jnp.asarray(gamma) + currents = jnp.asarray(currents) + + assert gamma.ndim == 3, "gamma must be a 3D array with shape (n_curves, n_segments, 3)" + assert gamma.shape[2] == 3, "gamma must have shape (n_curves, n_segments, 3)" + + if currents.ndim == 0: + currents = jnp.full((gamma.shape[0],), currents) + elif currents.ndim == 1 and currents.shape[0] == 1 and gamma.shape[0] != 1: + currents = jnp.full((gamma.shape[0],), currents[0]) + + assert isinstance(nfp, int) and nfp > 0, "nfp must be a positive integer" + assert isinstance(stellsym, bool), "stellsym must be a boolean" + assert currents.ndim == 1, "currents must be a scalar or a 1D array" + assert gamma.shape[0] == currents.shape[0], ( + f"Number of base curves must match number of base currents. " + f"Got gamma.shape[0]={gamma.shape[0]} and currents.shape[0]={currents.shape[0]}" + ) + + n_sym = nfp * (1 + int(stellsym)) + if n_sym > 1 and gamma.shape[0] % n_sym == 0: + n_base_candidate = gamma.shape[0] // n_sym + gamma_base_candidate = gamma[:n_base_candidate] + gamma_expanded_candidate = apply_symmetries_to_gammas(gamma_base_candidate, nfp, stellsym) + currents_base_candidate = currents[:n_base_candidate] + currents_expanded_candidate = apply_symmetries_to_currents(currents_base_candidate, nfp, stellsym) + + if ( + gamma_expanded_candidate.shape == gamma.shape + and currents_expanded_candidate.shape == currents.shape + and jnp.allclose(gamma_expanded_candidate, gamma) + and jnp.allclose(currents_expanded_candidate, currents) + ): + gamma = gamma_base_candidate + currents = currents_base_candidate + + self._gamma = gamma + self._dofs_currents_raw = currents + self._n_segments = gamma.shape[1] + self._nfp = nfp + self._stellsym = stellsym + self._scale_fixed = _initialize_scale_fixed(gamma, scale_fixed) + + self._gamma_dash = None + self._gamma_dashdash = None + self._length = None + self._curvature = None + self._currents_scale = _initialize_currents_scale(currents, currents_scale) + self._dofs_currents = None + self._currents = None + + # reset_cache method + def reset_cache(self): + self._gamma_dash = None + self._gamma_dashdash = None + self._length = None + self._curvature = None + self._dofs_currents = None + self._currents = None + + # dofs_gamma property and setter + @property + def dofs_gamma(self): + return jnp.array(self._gamma) / self.scale_fixed + + @dofs_gamma.setter + def dofs_gamma(self, new_dofs_gamma): + new_dofs_gamma = jnp.asarray(new_dofs_gamma) + assert new_dofs_gamma.ndim == 3, "dofs_gamma must have shape (n_base_curves, n_segments, 3)" + assert new_dofs_gamma.shape[2] == 3, "dofs_gamma must have shape (n_base_curves, n_segments, 3)" + self.reset_cache() + self._gamma = new_dofs_gamma * self.scale_fixed + self._n_segments = new_dofs_gamma.shape[1] + + # gamma property and setter (symmetry-expanded) + @property + def gamma(self): + return apply_symmetries_to_gammas(self._gamma, self.nfp, self.stellsym) + + @gamma.setter + def gamma(self, new_gamma): + new_gamma = jnp.asarray(new_gamma) + assert new_gamma.ndim == 3, "gamma must be a 3D array with shape (n_curves, n_segments, 3)" + assert new_gamma.shape[2] == 3, "gamma must have shape (n_curves, n_segments, 3)" + + n_sym = self.nfp * (1 + int(self.stellsym)) + n_base = self.n_base_curves + + if new_gamma.shape[0] == n_base: + self.reset_cache() + self._gamma = new_gamma + self._n_segments = new_gamma.shape[1] + return + assert new_gamma.shape[0] == n_base * n_sym, ( + f"Expected gamma with {n_base} (base) or {n_base*n_sym} (expanded) curves, " + f"got {new_gamma.shape[0]}" + ) + # Ordering in apply_symmetries_to_gammas ensures the first n_base curves are k=0, flip=False (base) + self.reset_cache() + self._gamma = new_gamma[:n_base] + self._n_segments = new_gamma.shape[1] + + # n_segments property + @property + def n_segments(self): + return self._n_segments + + @property + def n_base_curves(self): + return self.dofs_gamma.shape[0] + + # nfp property + @property + def nfp(self): + return self._nfp + + # stellsym property + @property + def stellsym(self): + return self._stellsym + + # scale_fixed property and setter + @property + def scale_fixed(self): + return self._scale_fixed + + @scale_fixed.setter + def scale_fixed(self, new_scale_fixed): + self._gamma = self.dofs_gamma * new_scale_fixed + self._scale_fixed = new_scale_fixed + self.reset_cache() + + # dofs_currents_raw property and setter + @property + def dofs_currents_raw(self): + return jnp.array(self._dofs_currents_raw) + + @dofs_currents_raw.setter + def dofs_currents_raw(self, new_dofs_currents_raw): + new_dofs_currents_raw = jnp.asarray(new_dofs_currents_raw) + assert new_dofs_currents_raw.ndim == 1, "dofs_currents_raw must be a 1D array" + assert new_dofs_currents_raw.shape[0] == self.n_base_curves, ( + f"Expected {self.n_base_curves} base currents, got {new_dofs_currents_raw.shape[0]}" + ) + self.reset_cache() + self._dofs_currents_raw = jnp.asarray(new_dofs_currents_raw) + + # currents_scale property and setter + @property + def currents_scale(self): + return self._currents_scale + + @currents_scale.setter + def currents_scale(self, new_currents_scale): + self._dofs_currents_raw = self.dofs_currents * new_currents_scale + self._currents_scale = new_currents_scale + self._currents = None + + # dofs_currents property and setter + @property + def dofs_currents(self): + if self._dofs_currents is None: + self._dofs_currents = self.dofs_currents_raw / self.currents_scale + return self._dofs_currents + + @dofs_currents.setter + def dofs_currents(self, new_dofs_currents): + self.dofs_currents_raw = new_dofs_currents * self.currents_scale + + # currents property + @property + def currents(self): + if self._currents is None: + self._currents = apply_symmetries_to_currents(self.dofs_currents_raw, self.nfp, self.stellsym) + return self._currents + + # dofs property and setter (flattened gamma + currents) + @property + def dofs(self): + return jnp.hstack([self.dofs_gamma.ravel(), self.dofs_currents]) + + @dofs.setter + def dofs(self, new_dofs): + n_gamma_dofs = jnp.size(self.dofs_gamma) + self.dofs_gamma = jnp.reshape(new_dofs[:n_gamma_dofs], self.dofs_gamma.shape) + self.dofs_currents = new_dofs[n_gamma_dofs:] + + # x property and setter (for compatibility with simsopt) + @property + def x(self): + return self.dofs + + @x.setter + def x(self, new_dofs): + self.dofs = new_dofs + + # Compute derivatives using finite differences (circular) + def _compute_gamma_dash(self): + """Compute first derivative using finite differences on periodic curve""" + base_gamma = self._gamma + gamma_shift_forward = jnp.roll(base_gamma, -1, axis=1) + gamma_shift_backward = jnp.roll(base_gamma, 1, axis=1) + base_gamma_dash = (gamma_shift_forward - gamma_shift_backward) / 2.0 * self._n_segments + return apply_symmetries_to_gammas(base_gamma_dash, self.nfp, self.stellsym) + + def _compute_gamma_dashdash(self): + """Compute second derivative using finite differences on periodic curve""" + base_gamma = self._gamma + gamma_shift_forward = jnp.roll(base_gamma, -1, axis=1) + gamma_shift_backward = jnp.roll(base_gamma, 1, axis=1) + base_gamma_dashdash = (gamma_shift_forward - 2.0 * base_gamma + gamma_shift_backward) * (self._n_segments ** 2) + return apply_symmetries_to_gammas(base_gamma_dashdash, self.nfp, self.stellsym) + + # gamma_dash property + @property + def gamma_dash(self): + if self._gamma_dash is None: + self._gamma_dash = self._compute_gamma_dash() + return self._gamma_dash + + # gamma_dashdash property + @property + def gamma_dashdash(self): + if self._gamma_dashdash is None: + self._gamma_dashdash = self._compute_gamma_dashdash() + return self._gamma_dashdash + + # length property + @property + def length(self): + if self._length is None: + self._length = jnp.mean(jnp.linalg.norm(self.gamma_dash, axis=2), axis=1) + return self._length + + # curvature property + @staticmethod + @jit + def compute_curvature(gammadash, gammadashdash): + return jnp.linalg.norm(jnp.cross(gammadash, gammadashdash, axis=1), axis=1) / jnp.linalg.norm(gammadash, axis=1)**3 + + @property + def curvature(self): + if self._curvature is None: + self._curvature = vmap(self.compute_curvature)(self.gamma_dash, self.gamma_dashdash) + return self._curvature + + # copy method + def copy(self): + coils = DiscretizedCoils(self.dofs_gamma.copy(), self.dofs_currents_raw.copy(), + nfp=self.nfp, stellsym=self.stellsym, + currents_scale=self.currents_scale, scale_fixed=self.scale_fixed) + + # Initialize caches + coils._gamma_dash = self._gamma_dash + coils._gamma_dashdash = self._gamma_dashdash + coils._length = self._length + coils._curvature = self._curvature + coils._dofs_currents = self.dofs_currents + coils._currents = self._currents + + return coils + + # magic methods + def __str__(self): + return f"DiscretizedCoils with {self.n_base_curves} base curves ({self.gamma.shape[0]} total)\n" \ + + f"n_segments: {self.n_segments}\n" \ + + f"nfp: {self.nfp}, stellsym: {self.stellsym}\n" \ + + f"Degrees of freedom shape: {self.dofs.shape}\n" \ + + f"Currents scaling factor: {self.currents_scale}\n" + + def __repr__(self): + return f"DiscretizedCoils with {self.n_base_curves} base curves ({self.gamma.shape[0]} total)\n" \ + + f"n_segments: {self.n_segments}\n" \ + + f"nfp: {self.nfp}, stellsym: {self.stellsym}\n" \ + + f"Degrees of freedom shape: {self.dofs.shape}\n" \ + + f"Currents scaling factor: {self.currents_scale}\n" + + def __len__(self): + return self.gamma.shape[0] + + def __getitem__(self, key): + if isinstance(key, int): + return DiscretizedCoils(jnp.expand_dims(self.gamma[key], 0), jnp.expand_dims(self.currents[key], 0), + nfp=1, stellsym=False, + currents_scale=self.currents_scale, scale_fixed=self.scale_fixed) + elif isinstance(key, (slice, jnp.ndarray)): + return DiscretizedCoils(self.gamma[key], self.currents[key], nfp=1, stellsym=False, + currents_scale=self.currents_scale, scale_fixed=self.scale_fixed) + else: + raise TypeError(f"Invalid argument type. Got {type(key)}, expected int, slice or jnp.ndarray.") + + def __add__(self, other): + if isinstance(other, DiscretizedCoils): + return DiscretizedCoils( + jnp.concatenate((self.gamma, other.gamma), axis=0), + jnp.concatenate((self.currents, other.currents), axis=0), + nfp=1, stellsym=False # Combined coils lose symmetry structure + ) + else: + raise TypeError(f"Invalid argument type. Got {type(other)}, expected DiscretizedCoils.") + + def __contains__(self, other): + if isinstance(other, DiscretizedCoils): + return jnp.all(jnp.isin(other.dofs, self.dofs)) + else: + raise TypeError(f"Invalid argument type. Got {type(other)}, expected DiscretizedCoils.") + + def __eq__(self, other): + if isinstance(other, DiscretizedCoils): + if self.dofs.shape != other.dofs.shape: + return False + return jnp.all(self.gamma == other.gamma) and jnp.all(self.dofs_currents == other.dofs_currents) + else: + raise TypeError(f"Invalid argument type. Got {type(other)}, expected DiscretizedCoils.") + + def __ne__(self, other): + return not self.__eq__(other) + + def __iter__(self): + self.iter_idx = 0 + return self + + def __next__(self): + if self.iter_idx < len(self): + result = self[self.iter_idx] + self.iter_idx += 1 + return result + else: + raise StopIteration + + # Saving and loading methods + def save_coils(self, filename: str, text=""): + """Save the coils to a file""" + with open(filename, "a") as file: + file.write(f"n_segments: {self.n_segments}\n") + file.write(f"nfp: {self.nfp}, stellsym: {self.stellsym}\n") + file.write(f"Base gamma dofs\n") + file.write(f"{repr(self.dofs_gamma.tolist())}\n") + file.write(f"Gamma fixed scaling\n") + file.write(f"{self.scale_fixed}\n") + file.write(f"Currents degrees of freedom\n") + file.write(f"{repr(self.dofs_currents.tolist())}\n") + file.write(f"Currents scaling factor\n") + file.write(f"{self.currents_scale}\n") + file.write(f"{text}\n") + + def to_json(self, filename: str): + """Save coils to JSON file""" + data = { + "n_segments": self.n_segments, + "nfp": self.nfp, + "stellsym": self.stellsym, + "dofs_gamma_raw": self._gamma.tolist(), + "dofs_currents": self.dofs_currents.tolist(), + "currents_scale": float(self.currents_scale), + "scale_fixed": float(self.scale_fixed), + } + import json + with open(filename, 'w') as file: + json.dump(data, file) + + @classmethod + def from_json(cls, filename: str): + """Create DiscretizedCoils from JSON file""" + import json + with open(filename, "r") as file: + data = json.load(file) + gamma_data = data.get("dofs_gamma_raw", data.get("dofs_gamma", data.get("gamma"))) + gamma = jnp.array(gamma_data) + currents_scale = data.get("currents_scale", None) + currents = jnp.array(data["dofs_currents"]) + if currents_scale is not None: + currents = currents * currents_scale + scale_fixed = data.get("scale_fixed", data.get("fixed_scale", None)) + if "dofs_gamma_raw" not in data and scale_fixed is not None: + gamma = gamma * scale_fixed + nfp = data.get("nfp", 1) + stellsym = data.get("stellsym", False) + if "dofs_gamma" not in data and gamma.shape[0] % (nfp * (1 + int(stellsym))) == 0: + n_base = gamma.shape[0] // (nfp * (1 + int(stellsym))) + gamma = gamma[:n_base] + currents = currents[:n_base] + return cls(gamma, currents, nfp=nfp, stellsym=stellsym, currents_scale=currents_scale, scale_fixed=scale_fixed) + + def plot(self, ax=None, show=True, plot_derivative=False, close=False, axis_equal=True, + color="brown", linewidth=3, label=None, **kwargs): + """Plot the coils""" + def rep(data): + if close: + return jnp.concatenate((data, [data[0]])) + else: + return data + import matplotlib.pyplot as plt + if ax is None or ax.name != "3d": + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + label_count = 0 + for gamma, gammadash in zip(self.gamma, self.gamma_dash): + x = rep(gamma[:, 0]) + y = rep(gamma[:, 1]) + z = rep(gamma[:, 2]) + if plot_derivative: + xt = rep(gammadash[:, 0]) + yt = rep(gammadash[:, 1]) + zt = rep(gammadash[:, 2]) + if label_count == 0: + ax.plot(x, y, z, **kwargs, color=color, linewidth=linewidth, label=label) + label_count += 1 + else: + ax.plot(x, y, z, **kwargs, color=color, linewidth=linewidth) + if plot_derivative: + ax.quiver(x, y, z, 0.1 * xt, 0.1 * yt, 0.1 * zt, arrow_length_ratio=0.1, color='r') + if axis_equal: + fix_matplotlib_3d(ax) + if show: + plt.show() + + def to_vtk(self, filename: str, close: bool = True, extra_data=None): + """Export coils to VTK format""" + try: + import numpy as np + except ImportError: + raise ImportError("The 'numpy' library is required. Please install it using 'pip install numpy'.") + try: + from pyevtk.hl import polyLinesToVTK + except ImportError: + raise ImportError("The 'pyevtk' library is required. Please install it using 'pip install pyevtk'.") + + def wrap(data): + return jnp.concatenate([data, jnp.array([data[0]])]) + + gammas = self.gamma + if close: + x = jnp.concatenate([wrap(gamma[:, 0]) for gamma in gammas]) + y = jnp.concatenate([wrap(gamma[:, 1]) for gamma in gammas]) + z = jnp.concatenate([wrap(gamma[:, 2]) for gamma in gammas]) + ppl = jnp.asarray([gamma.shape[0] + 1 for gamma in gammas]) + else: + x = jnp.concatenate([gamma[:, 0] for gamma in gammas]) + y = jnp.concatenate([gamma[:, 1] for gamma in gammas]) + z = jnp.concatenate([gamma[:, 2] for gamma in gammas]) + ppl = jnp.asarray([gamma.shape[0] for gamma in gammas]) + + data = jnp.concatenate([i * jnp.ones((ppl[i],)) for i in range(len(gammas))]) + pointData = {'idx': np.array(data)} + if extra_data is not None: + pointData = {**pointData, **extra_data} + polyLinesToVTK(str(filename), np.array(x), np.array(y), np.array(z), + pointsPerLine=np.array(ppl), pointData=pointData) + + def to_simsopt(self): + """Convert to simsopt coils""" + from simsopt.geo import CurveXYZFourier + from simsopt.field import coils_via_symmetries, Current as Current_SIMSOPT + + curves_simsopt = [] + currents_simsopt = [] + + # Fit Fourier coefficients from base gammas + for g, current in zip(self._gamma, self.dofs_currents_raw): + # Fit Fourier coefficients + order = (self.n_segments // 2) - 1 + dofs, _ = fit_dofs_from_coils(jnp.expand_dims(g, 0), order, self.n_segments) + + curve = CurveXYZFourier(self.n_segments, order) + curve.x = jnp.reshape(dofs[0], curve.x.shape) + curves_simsopt.append(curve) + currents_simsopt.append(Current_SIMSOPT(current)) + + return coils_via_symmetries(curves_simsopt, currents_simsopt, self.nfp, self.stellsym) + + @classmethod + def from_simsopt(cls, simsopt_coils, nfp: int = 1, stellsym: bool = False): + """Create from simsopt coils + + Args: + simsopt_coils: List of simsopt coils or path to simsopt file + nfp: Number of field periods (default: 1) + stellsym: Stellarator symmetry (default: False) + """ + if isinstance(simsopt_coils, str): + from simsopt import load + bs = load(simsopt_coils) + simsopt_coils = bs.coils + + gammas = [] + currents = [] + + for coil in simsopt_coils: + gamma = jnp.array(coil.curve.gamma()) + gammas.append(gamma) + currents.append(coil.current.get_value()) + + gamma_array = jnp.array(gammas) + currents_array = jnp.array(currents) + + n_sym = nfp * (1 + int(stellsym)) + if n_sym > 1 and gamma_array.shape[0] % n_sym == 0: + n_base = gamma_array.shape[0] // n_sym + gamma_array = gamma_array[:n_base] + currents_array = currents_array[:n_base] + + return cls(gamma_array, currents_array, nfp=nfp, stellsym=stellsym) + + @classmethod + def from_Coils(cls, coils: Coils): + """Create from a standard Coils object""" + base_gamma = Curves(coils.dofs_curves, coils.n_segments, nfp=1, stellsym=False).gamma + currents = coils.dofs_currents_raw + return cls(base_gamma, currents, nfp=coils.nfp, stellsym=coils.stellsym) + + def to_Coils(self, order: int = None) -> Coils: + """Convert to standard Coils object + + Args: + order: Fourier order for fitted curves (default: n_segments // 2 - 1) + """ + if order is None: + order = (self.n_segments // 2) - 1 + + dofs, _ = fit_dofs_from_coils(self._gamma, order, self.n_segments) + curves = Curves(dofs, self.n_segments, nfp=self.nfp, stellsym=self.stellsym) + return Coils(curves, self.dofs_currents_raw) + + def _tree_flatten(self): + children = (self.dofs_gamma, self.dofs_currents) + aux_data = { + "n_segments": self._n_segments, + "nfp": self._nfp, + "stellsym": self._stellsym, + "currents_scale": self.currents_scale, + "scale_fixed": self.scale_fixed, + } + return (children, aux_data) + + @classmethod + def _tree_unflatten(cls, aux_data, children): + dofs_gamma, dofs_currents = children + return cls( + dofs_gamma * aux_data["scale_fixed"], + dofs_currents * aux_data["currents_scale"], + nfp=aux_data["nfp"], + stellsym=aux_data["stellsym"], + currents_scale=aux_data["currents_scale"], + scale_fixed=aux_data["scale_fixed"], + ) + +tree_util.register_pytree_node(DiscretizedCoils, + DiscretizedCoils._tree_flatten, + DiscretizedCoils._tree_unflatten) diff --git a/essos/objective_functions.py b/essos/objective_functions.py index 16257d52..08dd3307 100644 --- a/essos/objective_functions.py +++ b/essos/objective_functions.py @@ -11,7 +11,7 @@ from essos.coils import Curves, Coils from essos.optimization import new_nearaxis_from_x_and_old_nearaxis from essos.constants import mu_0 -from essos.coil_perturbation import perturb_curves_systematic, perturb_curves_statistic +from essos.coil_perturbation import perturb_curves @@ -29,8 +29,8 @@ def perturbed_coils_from_dofs(x,key,sampler,dofs_curves,currents_scale,nfp,n_seg #Split once the key/seed given for one pertubred stellarator split_keys = jax.random.split(jax.random.key(key), 2) #Internally the following functions will then further split the two keys avoiding repeating keys - perturb_curves_systematic(coils.curves, sampler, key=split_keys[0]) - perturb_curves_statistic(coils.curves, sampler, key=split_keys[1]) + coils = perturb_curves(coils, sampler, key=split_keys[0], perturbation_type='systematic') + coils = perturb_curves(coils, sampler, key=split_keys[1], perturbation_type='statistical') return coils def field_from_dofs(x,dofs_curves,currents_scale,nfp,n_segments=60, stellsym=True): diff --git a/essos/surfaces.py b/essos/surfaces.py index 78e5cc02..63df5ff6 100644 --- a/essos/surfaces.py +++ b/essos/surfaces.py @@ -109,8 +109,27 @@ def nested_lists_to_array(ll): class SurfaceRZFourier: def __init__(self, rc, zs, nfp, mpol, ntor, ntheta=30, nphi=30, close=True, range_torus='full torus', scaling_type=2, scaling_factor=0): - """ rc, zs: dynamic arrays - nfp, mpol, ntor: static """ + """Initialize a Fourier surface. + + Args: + rc: cosine Fourier coefficients for R. + zs: sine Fourier coefficients for Z. + nfp: number of field periods. + mpol: maximum poloidal mode number. + ntor: maximum toroidal mode number. + ntheta: number of theta grid points. + nphi: number of phi grid points. + close: whether the surface mesh includes the endpoint. + range_torus: either ``'full torus'`` or ``'half period'``. + scaling_type: norm used in the mode scaling. Accepted values are + ``'L1'`` or ``1``, ``'L2'`` or ``2``, and ``'Linfty'`` or ``-1``. + scaling_factor: exponential weight used in the scaling + ``exp(scaling_factor * ||(xm, xn)||)``. + + Note: + The optimized dofs are stored as ``[rc * scaling, zs * scaling]``, + with the scaling computed mode-by-mode from ``xm`` and ``xn``. + """ assert isinstance(nfp, int) and nfp > 0, "nfp must be a positive integer." assert isinstance(mpol, int) and mpol >= 0, "mpol must be a non-negative integer." @@ -146,10 +165,24 @@ def __init__(self, rc, zs, nfp, mpol, ntor, ntheta=30, nphi=30, close=True, rang self._phi2d = None self._angles = None - self._scaling_type = scaling_type # 1 for L-1 norm, 2 for L-2 norm, jnp.inf for L-infinity norm + self._scaling_type = self._normalize_scaling_type(scaling_type) self._scaling_factor = scaling_factor self._scaling = None + @staticmethod + def _normalize_scaling_type(scaling_type): + """Map public scaling_type inputs to norm orders used internally.""" + if scaling_type == "L1" or scaling_type == 1: + return 1 + if scaling_type == "L2" or scaling_type == 2: + return 2 + if scaling_type == "Linfty" or scaling_type == -1 or scaling_type == jnp.inf: + return jnp.inf + raise ValueError( + f"Unknown scaling_type: {scaling_type}. " + "Expected 'L1', 1, 'L2', 2, 'Linfty', -1, or jnp.inf." + ) + @classmethod def from_input_file(cls, file, ntheta=30, nphi=30, close=True, range_torus='full torus'): @@ -353,7 +386,7 @@ def scaling_type(self): @scaling_type.setter def scaling_type(self, new_type): - self._scaling_type = new_type + self._scaling_type = self._normalize_scaling_type(new_type) self._scaling = None # scaling_factor property and setter @@ -369,6 +402,7 @@ def scaling_factor(self, new_factor): # scaling property @property def scaling(self): + """Mode-by-mode scaling ``exp(scaling_factor * ||(xm, xn)||)``.""" if self._scaling is None: self._scaling = jnp.exp(self.scaling_factor * jnp.linalg.norm(jnp.vstack([self.xm, self.xn]), ord=self.scaling_type, axis=0)) return self._scaling @@ -655,7 +689,7 @@ def mean_cross_sectional_area(self): return mean_cross_sectional_area def _tree_flatten(self): - children = (self._rc, self._zs) # arrays / dynamic values + children = (self.dofs,) # arrays / dynamic values aux_data = {"nfp": self._nfp, "mpol": self._mpol, "ntor": self._ntor, @@ -669,7 +703,24 @@ def _tree_flatten(self): @classmethod def _tree_unflatten(cls, aux_data, children): - return cls(*children, **aux_data) + dofs, = children + half = dofs.size // 2 + rc_scaled = dofs[:half] + zs_scaled = dofs[half:] + + mpol = aux_data["mpol"] + ntor = aux_data["ntor"] + nfp = aux_data["nfp"] + scaling_type = cls._normalize_scaling_type(aux_data["scaling_type"]) + scaling_factor = aux_data["scaling_factor"] + + xm = jnp.repeat(jnp.arange(mpol + 1), 2 * ntor + 1)[ntor:] + xn = nfp * jnp.tile(jnp.arange(-ntor, ntor + 1), mpol + 1)[ntor:] + scaling = jnp.exp(scaling_factor * jnp.linalg.norm(jnp.vstack([xm, xn]), ord=scaling_type, axis=0)) + + rc = rc_scaled / scaling + zs = zs_scaled / scaling + return cls(rc, zs, **aux_data) tree_util.register_pytree_node(SurfaceRZFourier, SurfaceRZFourier._tree_flatten, @@ -805,6 +856,3 @@ def plot_scalar_on_flux_surface(surface, scalar_map): surface: the surface object in which to plot the scalar_map scalar_map: a scalar_map as function of theta and phi ''' - - - diff --git a/examples/simple_examples/create_perturbed_coils.py b/examples/simple_examples/create_perturbed_coils.py index 271181f6..63683667 100644 --- a/examples/simple_examples/create_perturbed_coils.py +++ b/examples/simple_examples/create_perturbed_coils.py @@ -10,8 +10,7 @@ import matplotlib.pyplot as plt from essos.coils import Coils, CreateEquallySpacedCurves,Curves from functools import partial -from essos.coil_perturbation import GaussianSampler -from essos.coil_perturbation import perturb_curves_statistic,perturb_curves_systematic +from essos.coil_perturbation import GaussianSampler, perturb_curves @@ -42,14 +41,14 @@ split_keys=jax.random.split(jax.random.key(key), num=2) #Add systematic error coils_sys = Coils(curves=curves, currents=[current_on_each_coil]*number_coils_per_half_field_period) -perturb_curves_systematic(coils_sys.curves, g, key=split_keys[0]) +coils_sys = perturb_curves(coils_sys, g, key=split_keys[0], perturbation_type='systematic') # Add statistical error coils_stat = Coils(curves=curves, currents=[current_on_each_coil]*number_coils_per_half_field_period) -perturb_curves_statistic(coils_stat.curves, g, key=split_keys[1]) +coils_stat = perturb_curves(coils_stat, g, key=split_keys[1], perturbation_type='statistical') # Add both systematic and statistical errors coils_perturbed = Coils(curves=curves, currents=[current_on_each_coil]*number_coils_per_half_field_period) -perturb_curves_systematic(coils_perturbed.curves, g, key=split_keys[0]) -perturb_curves_statistic(coils_perturbed.curves, g, key=split_keys[1]) +coils_perturbed = perturb_curves(coils_perturbed, g, key=split_keys[0], perturbation_type='systematic') +coils_perturbed = perturb_curves(coils_perturbed, g, key=split_keys[1], perturbation_type='statistical') fig = plt.figure(figsize=(9, 8)) diff --git a/tests/test_coils.py b/tests/test_coils.py index c9e1f58d..afef5f19 100644 --- a/tests/test_coils.py +++ b/tests/test_coils.py @@ -1,4 +1,5 @@ import pytest +import jax from essos.coils import Curves import jax.numpy as jnp import random @@ -47,6 +48,15 @@ def test_curves_property_setters(): curves.stellsym = False assert curves.stellsym == False +def test_curves_pytree_preserves_scaling_metadata(): + dofs = jnp.ones((2, 3, 5)) + curves = Curves(dofs, scaling_type=2, scaling_factor=0.3, scale_fixed=7.0) + curves_copy = jax.tree_util.tree_map(lambda x: x, curves) + + assert curves_copy.scaling_type == curves.scaling_type + assert curves_copy.scaling_factor == curves.scaling_factor + assert curves_copy.scale_fixed == curves.scale_fixed + def test_curves_str_repr(): dofs = jnp.zeros((2, 3, 5)) curves = Curves(dofs) @@ -116,4 +126,4 @@ def test_curves_iter(): assert curve.curves.shape == (1, 3, 5) if __name__ == "__main__": - pytest.main() \ No newline at end of file + pytest.main() diff --git a/tests/test_objective_functions.py b/tests/test_objective_functions.py index c2d6eedc..857e8feb 100644 --- a/tests/test_objective_functions.py +++ b/tests/test_objective_functions.py @@ -92,9 +92,8 @@ def setUp(self): @patch('essos.objective_functions.Curves', return_value=DummyCurves()) @patch('essos.objective_functions.Coils', return_value=DummyCoils()) @patch('essos.objective_functions.BiotSavart', return_value=DummyField()) - @patch('essos.objective_functions.perturb_curves_systematic') - @patch('essos.objective_functions.perturb_curves_statistic') - def test_perturbed_field_and_coils_from_dofs(self, pcs, pcss, bs, coils, curves): + @patch('essos.objective_functions.perturb_curves', side_effect=lambda curves, sampler, key=None, perturbation_type=None: curves) + def test_perturbed_field_and_coils_from_dofs(self, pc, bs, coils, curves): objf.pertubred_field_from_dofs(self.x, self.key, self.sampler, self.dofs_curves, self.currents_scale, self.nfp) objf.perturbed_coils_from_dofs(self.x, self.key, self.sampler, self.dofs_curves, self.currents_scale, self.nfp) @@ -258,4 +257,4 @@ def test_linking_number_pure_and_integrand(self): objf.integrand_linking_number(r1, dr1, r2, dr2, dphi, dphi) if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main()