diff --git a/debug_from_simsopt_import.py b/debug_from_simsopt_import.py new file mode 100644 index 00000000..bff6816c --- /dev/null +++ b/debug_from_simsopt_import.py @@ -0,0 +1,96 @@ +import os +import jax +import jax.numpy as jnp + +from essos.coils import Coils +from essos.fields import BiotSavart, Vmec +from essos.constants import ( + ALPHA_PARTICLE_MASS, + ALPHA_PARTICLE_CHARGE, + FUSION_ALPHA_PARTICLE_ENERGY, +) +from essos.dynamics import GuidingCenter, Particles +from essos.objective_functions import normB_axis + + +def describe_case(label, coils, vmec_surface, initial_xyz, initial_vpar): + print(f"\n=== {label} ===") + print("n_base_curves:", coils.curves.n_base_curves) + print("n_segments:", coils.n_segments) + print("nfp:", coils.nfp, "stellsym:", coils.stellsym) + print("dofs_curves shape:", coils.dofs_curves.shape) + print("dofs_currents_raw shape:", coils.dofs_currents_raw.shape) + print("currents raw min/max:", float(jnp.min(coils.dofs_currents_raw)), float(jnp.max(coils.dofs_currents_raw))) + print("currents scale:", float(coils.currents_scale)) + print("currents normalized min/max:", float(jnp.min(coils.dofs_currents)), float(jnp.max(coils.dofs_currents))) + print("gamma shape:", coils.gamma.shape) + print("gamma finite:", bool(jnp.all(jnp.isfinite(coils.gamma)))) + print("gamma_dash finite:", bool(jnp.all(jnp.isfinite(coils.gamma_dash)))) + print("gamma_dashdash finite:", bool(jnp.all(jnp.isfinite(coils.gamma_dashdash)))) + + field = BiotSavart(coils) + B_axis = normB_axis(field, npoints=200) + print("axis B mean before renorm:", float(jnp.mean(B_axis))) + coils.dofs_currents = coils.dofs_currents * 5.7 / jnp.mean(B_axis) + field = BiotSavart(coils) + print("axis B mean after renorm:", float(jnp.mean(normB_axis(field, npoints=200)))) + + print("surface gamma finite:", bool(jnp.all(jnp.isfinite(vmec_surface.gamma)))) + print("initial xyz:", initial_xyz) + print("field.B(initial):", field.B(initial_xyz)) + print("field.AbsB(initial):", float(field.AbsB(initial_xyz))) + print("field.dAbsB_by_dX(initial):", field.dAbsB_by_dX(initial_xyz)) + print("field.curl_b(initial):", field.curl_b(initial_xyz)) + print("field.kappa(initial):", field.kappa(initial_xyz)) + + finite_checks = { + "B": jnp.all(jnp.isfinite(field.B(initial_xyz))), + "AbsB": jnp.isfinite(field.AbsB(initial_xyz)), + "dAbsB_by_dX": jnp.all(jnp.isfinite(field.dAbsB_by_dX(initial_xyz))), + "curl_b": jnp.all(jnp.isfinite(field.curl_b(initial_xyz))), + "kappa": jnp.all(jnp.isfinite(field.kappa(initial_xyz))), + } + print("finite checks:", {k: bool(v) for k, v in finite_checks.items()}) + + particles = Particles( + initial_xyz=jnp.expand_dims(initial_xyz, 0), + mass=ALPHA_PARTICLE_MASS, + charge=ALPHA_PARTICLE_CHARGE, + energy=FUSION_ALPHA_PARTICLE_ENERGY, + ) + initial_condition = jnp.array([initial_xyz[0], initial_xyz[1], initial_xyz[2], initial_vpar]) + rhs = GuidingCenter(0.0, initial_condition, (field, particles, type("E", (), {"E_covariant": lambda self, x: jnp.zeros(3)})())) + print("GuidingCenter rhs:", rhs) + print("GuidingCenter rhs finite:", bool(jnp.all(jnp.isfinite(rhs)))) + + +def main(): + print(jax.devices()) + + simsopt_json = os.path.join("examples", "input_files", "QH_simple_scaled.json") + wout_file = os.path.join("examples", "input_files", "wout_QH_simple_scaled.nc") + vmec = Vmec(wout_file) + + R0 = 17.0 + initial_xyz = jnp.array([R0, 0.0, 0.0]) + initial_vpar = 0.0 + + coils_true = Coils.from_simsopt(simsopt_json, nfp=4, stellsym=True) + describe_case("from_simsopt stellsym=True", coils_true, vmec.surface, initial_xyz, initial_vpar) + + coils_false = Coils.from_simsopt(simsopt_json, nfp=4, stellsym=False) + describe_case("from_simsopt stellsym=False", coils_false, vmec.surface, initial_xyz, initial_vpar) + + from simsopt import load + + bs = load(simsopt_json) + base_coils = bs.coils[:6] + print("\n=== direct simsopt base coils ===") + print("n base coils:", len(base_coils)) + print("current values:", [float(c.current.get_value()) for c in base_coils]) + print("curve gamma shape:", jnp.asarray(base_coils[0].curve.gamma()).shape) + print("curve gamma first point:", jnp.asarray(base_coils[0].curve.gamma())[0]) + + +if __name__ == "__main__": + main() diff --git a/essos/__main__.py b/essos/__main__.py index 4e22aeb2..69faf914 100644 --- a/essos/__main__.py +++ b/essos/__main__.py @@ -6,7 +6,9 @@ import jax.numpy as jnp import matplotlib.pyplot as plt from essos.coils import Coils, CreateEquallySpacedCurves -from essos.fields import near_axis, BiotSavart +from essos.fields import BiotSavart +from pyqsc_jax.near_axis import near_axis + from essos.dynamics import Tracing from essos.optimization import optimize_loss_function from essos.objective_functions import loss_coils_for_nearaxis diff --git a/essos/augmented_lagrangian.py b/essos/augmented_lagrangian.py index 222c9c51..f68c500b 100644 --- a/essos/augmented_lagrangian.py +++ b/essos/augmented_lagrangian.py @@ -15,150 +15,312 @@ class LagrangeMultiplier(NamedTuple): """A class containing constrain parameters for Augmented Lagrangian Method""" value: Any penalty: Any + omega: Any + eta: Any sq_grad: Any #For updating squared gradient in case of adaptative penalty and multiplier evolution +def _multiplier_like(out, multiplier, penalty, omega, eta, sq_grad): + if out is None: + raise ValueError( + "Constraint function returned None during initialization. " + "Constraints used with eq()/ineq() must return a scalar or array." + ) + z = jnp.zeros_like(out) + return LagrangeMultiplier( + value=multiplier + z, + penalty=penalty + z, + omega=omega + z, + eta=eta + z, + sq_grad=sq_grad + z, + ) -#This is used for the usual augmented lagrangian form -def update_method(params,updates,eta,omega,model_mu='Constant',beta=2.0,mu_max=1.e4,alpha=0.99,gamma=1.e-2,epsilon=1.e-8,eta_tol=1.e-4,omega_tol=1.e-6): - """Different methods for updating multipliers and penalties - """ - pred = lambda x: isinstance(x, LagrangeMultiplier) - if model_mu=='Constant': - #jax.debug.print('{m}', m=model_mu) - return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(y.value,0.0*x.value,0.0*x.value),params,updates,is_leaf=pred) - elif model_mu=='Mu_Monotonic': - #jax.debug.print('{m}', m=model_mu) - return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred) - elif model_mu=='Mu_Conditional_True': - #jax.debug.print('True {m}', m=model_mu) - return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,0.0*x.value,0.0*x.value),params,updates,is_leaf=pred) - elif model_mu=='Mu_Conditional_False': - #jax.debug.print('False {m}', m=model_mu) - return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred) - elif model_mu=='Mu_Tolerance_True': - #jax.debug.print('Standard True {m}', m=model_mu) - mu_average=penalty_average(params) - #eta=eta/mu_average**(0.1) - #omega=omega/mu_average - eta=jnp.maximum(eta/mu_average**(0.1),eta_tol) - omega=jnp.maximum(omega/mu_average,omega_tol) - return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,0.0*x.value,0.0*x.value),params,updates,is_leaf=pred),eta,omega - elif model_mu=='Mu_Tolerance_False': - #jax.debug.print('Standard False {m}', m=model_mu) - mu_average=penalty_average(params) - #eta=1./mu_average**(0.1) - #omega=1./mu_average - eta=jnp.maximum(1./mu_average**(0.1),eta_tol) - #jax.debug.print('HMMMMMM mu_av {m}', m=mu_average) - #jax.debug.print('HMMMMMM eta {m}', m=eta) - omega=jnp.maximum(1./mu_average,omega_tol) - return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega - #return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega - elif model_mu=='Mu_Adaptative': - #jax.debug.print('True {m}', m=model_mu) - #Note that y.penalty is the derivative with respect to mu and so it is 0.5*C(x)**2, like the derivative with respect to lambda is C(x) - return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(gamma/(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon)*y.value,-x.penalty+gamma/(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon),-x.sq_grad+alpha*x.sq_grad+(1.-alpha)*y.penalty*2.),params,updates,is_leaf=pred) - - - -#This is used for the squared form of the augmented Lagrangioan -def update_method_squared(params,updates,eta,omega,model_mu='Constant',beta=2.0,mu_max=1.e4,alpha=0.99,gamma=1.e-2,epsilon=1.e-8,eta_tol=1.e-4,omega_tol=1.e-6): - """Different methods for updating multipliers and penalties) +class BaseConstraint: + """A minimal mutable container holding `init` and `loss` callables for a constraint. + + This mirrors the simple tuple-like behavior used elsewhere but allows + attribute access and matches the `base_loss` style in `losses.py`. """ + def __init__(self, init: Callable, loss: Callable): + self.init = init + self.loss = loss - pred = lambda x: isinstance(x, LagrangeMultiplier) - if model_mu=='Constant': - #jax.debug.print('{m}', m=model_mu) - return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier((y.value-x.value/x.penalty),0.0*x.value,0.0*x.value),params,updates,is_leaf=pred) - elif model_mu=='Mu_Monotonic': - #jax.debug.print('{m}', m=model_mu) - return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*(y.value-x.value/x.penalty),-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred) - elif model_mu=='Mu_Conditional_True': - #jax.debug.print('True {m}', m=model_mu) - return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*(y.value-x.value/x.penalty),0.0*x.value,0.0*x.value),params,updates,is_leaf=pred) - elif model_mu=='Mu_Conditional_False': - #jax.debug.print('False {m}', m=model_mu) - return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred) - elif model_mu=='Mu_Tolerance_True': - #jax.debug.print('Squared True {m}', m=model_mu) - mu_average=penalty_average(params) - #eta=eta/mu_average**(0.1) - #omega=omega/mu_average - eta=jnp.maximum(eta/mu_average**(0.1),eta_tol) - omega=jnp.maximum(omega/mu_average,omega_tol) - return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*(y.value-x.value/x.penalty),0.0*x.value,0.0*x.value),params,updates,is_leaf=pred),eta,omega - elif model_mu=='Mu_Tolerance_False': - #jax.debug.print('Squared False {m}', m=model_mu) - mu_average=penalty_average(params) - #eta=1./mu_average**(0.1) - #omega=1./mu_average - eta=jnp.maximum(1./mu_average**(0.1),eta_tol) - omega=jnp.maximum(1./mu_average,omega_tol) - return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega - #return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega - elif model_mu=='Mu_Adaptative': - #jax.debug.print('True {m}', m=model_mu) - #Note that y.penalty is the derivative with respect to mu and so it is 0.5*C(x)**2, like the derivative with respect to lambda is C(x) - return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(gamma/(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon)*(y.value-x.value/x.penalty),-x.penalty+gamma/(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon),-x.sq_grad+alpha*x.sq_grad+(1.-alpha)*(y.penalty*2.+(x.value/x.penalty)**2)),params,updates,is_leaf=pred) - - - - -def lagrange_update(model_lagrangian='Standard'): - """A gradient transformation for Optax that prepares an MDMM gradient - descent ascent update from a normal gradient descent update. - - It should be used like this with a base optimizer: - optimizer = optax.chain( - optax.sgd(1e-3), - mdmm_jax.optax_prepare_update(), - ) +class CompositeConstraint: + """Mutable composite constraint container. - Returns: - An Optax gradient transformation that converts a gradient descent update - into a gradient descent ascent update. + Exposes `init` and `loss` callables (same as `Constraint`) while + allowing attaching metadata like `arg_names`, `_dependencies`, and + `selective_map`. `set_dependencies` will propagate dependencies to any + contained `SelectiveConstraint` instances. """ - def init_fn(params): - del params - return optax.EmptyState() - - def update_fn(lagrange_params,updates, state,eta,omega, params=None,model_mu='Constant',beta=2.,mu_max=1.e4,alpha=0.99,gamma=1.e-2,epsilon=1.e-8,eta_tol=1.e-4,omega_tol=1.e-6): - del params - if model_lagrangian=='Standard' : - return update_method(lagrange_params,updates,eta,omega,model_mu=model_mu,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol), state - elif model_lagrangian=='Squared' : - return update_method_squared(lagrange_params,updates,eta,omega,model_mu=model_mu,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol), state + def __init__(self, init_fn: Callable, loss_fn: Callable, selective_map=None, arg_names=None): + self.init = init_fn + self.loss = loss_fn + self.selective_map = selective_map or {} + # keep ordered list of dependency names + self.arg_names = list(arg_names) if arg_names is not None else [] + self._dependencies = {} + self._starting_dofs = None + self._dofs_to_pytree = None + + def clear_cache(self): + self._starting_dofs = None + self._dofs_to_pytree = None + + @property + def dependencies(self): + return self._dependencies + + @dependencies.setter + def dependencies(self, value): + if not isinstance(value, dict): + raise TypeError("dependencies must be a dictionary mapping names to arrays") + self.clear_cache() + self._dependencies = value + for selective in self.selective_map.values(): + selective.dependencies = value + + def set_dependencies(self, deps): + self.dependencies = deps + + @property + def starting_dofs(self): + if self._starting_dofs is None: + if not self._dependencies: + raise RuntimeError("dependencies must be set on composite before accessing starting_dofs") + vals = tuple(self._dependencies[name] for name in self.arg_names) + self._starting_dofs, self._dofs_to_pytree = jax.flatten_util.ravel_pytree(vals) + return self._starting_dofs + + @property + def dofs_to_pytree(self): + if self._dofs_to_pytree is None: + _ = self.starting_dofs + return self._dofs_to_pytree + + +class SelectiveConstraint: + """Wraps a constraint with selective named dependencies, similar to custom_loss. + + This allows constraints to only depend on a subset of the available degrees of freedom + by name, enabling combination of constraints with different argument requirements. + No need to specify indices - just the dependency names! + + Why filtering is necessary: + - Different constraints may require different subsets of arguments (e.g., one depends on + 'field' only, another on 'coil' only, a third on both). + - Named filtering ensures each constraint receives only its required arguments, avoiding: + * Unnecessary computations with unused data + * Constraints expecting different signatures from breaking + * Wasteful memory transfer of irrelevant arrays + - Enables flexible constraint composition where DOF dependencies vary. + + Attributes: + constraint: The underlying Constraint (init_fn, loss_fn) tuple + arg_names: Tuple of argument names that this constraint depends on + dependencies: Dictionary mapping dependency names to their arrays/objects + + Example: + # Create constraints on different DOF subsets - no indices needed! + field_constraint = alm.eq(lambda field: jnp.sum(field**2)) + surface_constraint = alm.eq(lambda surface: jnp.sum(surface**2)) + + selective1 = SelectiveConstraint(field_constraint, 'field') + selective2 = SelectiveConstraint(surface_constraint, 'surface') + + combined = alm.combine(selective1, selective2) + # Set dependencies by name + combined.dependencies = {'field': field_array, 'surface': surface_array} + """ + def __init__(self, constraint: BaseConstraint, *arg_names, **kwargs): + """Initialize a SelectiveConstraint with named dependencies. + + Args: + constraint: A Constraint (init_fn, loss_fn) tuple + *arg_names: Names of the arguments this constraint depends on (order matters!) + """ + if not (hasattr(constraint, 'init') and hasattr(constraint, 'loss')): + raise TypeError(f"constraint must provide `init` and `loss` callables, got {type(constraint)}") + if not arg_names: + raise ValueError("At least one argument name must be provided") + + self.constraint = constraint + self.arg_names = arg_names + self.kwargs = kwargs + self._dependencies = {} + self._starting_dofs = None + self._dofs_to_pytree = None + + def clear_cache(self): + self._starting_dofs = None + self._dofs_to_pytree = None + + @property + def dependencies(self): + """Get the dependencies dictionary.""" + return self._dependencies + + @dependencies.setter + def dependencies(self, value): + """Set dependencies (mapping of arg_names to their values).""" + if not isinstance(value, dict): + raise TypeError("dependencies must be a dictionary mapping names to arrays") + self.clear_cache() + self._dependencies = value + + def _get_filtered_args(self): + """Extract only the required arguments from dependencies by name.""" + filtered_args = [] + for name in self.arg_names: + if name not in self._dependencies: + raise KeyError( + f"SelectiveConstraint '{self.arg_names}' depends on '{name}', " + f"but it's not in dependencies dict. Available: {list(self._dependencies.keys())}" + ) + filtered_args.append(self._dependencies[name]) + return tuple(filtered_args) + + @property + def starting_dofs(self): + if self._starting_dofs is None: + if not self._dependencies: + raise RuntimeError("dependencies must be set before accessing starting_dofs") + vals = tuple(self._dependencies[name] for name in self.arg_names) + self._starting_dofs, self._dofs_to_pytree = jax.flatten_util.ravel_pytree(vals) + return self._starting_dofs + + @property + def dofs_to_pytree(self): + if self._dofs_to_pytree is None: + _ = self.starting_dofs + return self._dofs_to_pytree + + def init(self, *args, **kwargs): + """Initialize constraint parameters using current dependencies.""" + # Prefer explicit call-time args/kwargs; otherwise use stored dependencies and constructor kwargs + if args or kwargs: + return self.constraint.init(*args, **kwargs, **self.kwargs) + args = self._get_filtered_args() + return self.constraint.init(*args, **self.kwargs) + + def loss(self, params, *args, **kwargs): + """Compute loss using current dependencies. + + Args: + params: Constraint parameters from init() + + Returns: + (loss_value, constraint_info) tuple + """ + # Prefer explicit call-time args/kwargs; otherwise use stored dependencies and constructor kwargs + if args or kwargs: + return self.constraint.loss(params, *args, **kwargs, **self.kwargs) + args = self._get_filtered_args() + return self.constraint.loss(params, *args, **self.kwargs) + + +class ScaledConstraint: + """Wraps a constraint with automatic or user-specified scaling. + + Useful when combining constraints with vastly different magnitudes. + Automatically normalizes constraint output by its initial norm or element-wise absolute values. + + Attributes: + constraint: The underlying Constraint (init_fn, loss_fn) tuple + scale_factor: Scalar to multiply constraint output (auto-computed or user-specified) + auto_scale: Whether to auto-compute scale_factor from initial constraint norm + elementwise_scale: If True, scale by 1/abs(constraint) element-wise; if False, scale by 1/norm + + Example: + # Auto-scale based on initial constraint norm (default) + constraint = alm.eq(lambda x: jnp.sum(x**2)) + scaled = ScaledConstraint(constraint, auto_scale=True, elementwise_scale=False) + + # Auto-scale element-wise + scaled = ScaledConstraint(constraint, auto_scale=True, elementwise_scale=True) + + # Or use fixed scaling + scaled = ScaledConstraint(constraint, scale_factor=1.0, auto_scale=False) + """ + def __init__(self, constraint: BaseConstraint, scale_factor=1.0, auto_scale=True, elementwise_scale=False): + """Initialize a ScaledConstraint. + + Args: + constraint: A Constraint (init_fn, loss_fn) tuple + scale_factor: Fixed scaling factor (ignored if auto_scale=True) + auto_scale: If True, automatically compute scale_factor from initial constraint + elementwise_scale: If True, scale element-wise by 1/abs(constraint); if False, scale by 1/norm + """ + if not (hasattr(constraint, 'init') and hasattr(constraint, 'loss')): + raise TypeError(f"constraint must provide `init` and `loss` callables, got {type(constraint)}") + + if not auto_scale and scale_factor <= 0: + raise ValueError(f"scale_factor must be positive, got {scale_factor}") + + self.constraint = constraint + self.scale_factor = scale_factor + self.auto_scale = auto_scale + self.elementwise_scale = elementwise_scale + self._initial_scale = None + self._initial_scale_scalar = None # For loss_value scaling + + def init(self, *args, **kwargs): + """Initialize constraint parameters and compute scale factor if needed.""" + params = self.constraint.init(*args, **kwargs) + + # If auto_scale, compute initial constraint scaling + if self.auto_scale: + _, constraint_info = self.constraint.loss(params, *args, **kwargs) + constraint_flat, unflatten_fn = jax.flatten_util.ravel_pytree(constraint_info) + + if self.elementwise_scale: + # Element-wise scaling: scale = 1 / abs(constraint_flat) + scale_flat = 1.0 / jnp.maximum(jnp.abs(constraint_flat), 1.) + # Unflatten to match constraint_info structure + self._initial_scale = unflatten_fn(scale_flat) + # For loss_value, use mean of element-wise scales + self._initial_scale_scalar = jnp.linalg.norm(scale_flat) + else: + # Norm-based scaling: scale = 1 / norm(constraint_flat) + scale_scalar = 1.0 / jnp.maximum(jnp.linalg.norm(constraint_flat), 1.) + self._initial_scale = scale_scalar + self._initial_scale_scalar = scale_scalar + + return params + + def loss(self, params, *args, **kwargs): + """Compute scaled constraint loss. + + Args: + params: Constraint parameters from init() + *args: Arguments to constraint + **kwargs: Keyword arguments + + Returns: + (scaled_loss_value, constraint_info) tuple + """ + loss_value, constraint_info = self.constraint.loss(params, *args, **kwargs) + + # Apply scaling + if self.auto_scale and self._initial_scale is not None: + # Scale loss_value with scalar + loss_scale = self._initial_scale_scalar + # Scale constraint_info with potentially element-wise scale + info_scale = self._initial_scale else: - print('Lagrangian model not available please select Standard or Squared ') - os._exit(0) - - return optax.GradientTransformation(init_fn, update_fn) - + loss_scale = self.scale_factor + info_scale = self.scale_factor + + return loss_value * loss_scale, jax.tree_util.tree_map(lambda x: x * info_scale, constraint_info) +def eq(fun,model_lagrangian='Standard', multiplier=0.0,penalty=1.,omega=1.0,eta=1.0,sq_grad=0., weight=1., reduction=jnp.sum): -class Constraint(NamedTuple): - """A pair of pure functions implementing a constraint. - - Attributes: - init: A pure function which, when called with an example instance of - the arguments to the constraint functions, returns a pytree - containing the constraint's learnable parameters. - loss: A pure function which, when called with the the learnable - parameters returned by init() followed by the arguments to the - constraint functions, returns the loss value for the constraint. - """ - init: Callable - loss: Callable - - -def eq(fun,model_lagrangian='Standard', multiplier=0.0,penalty=1.,sq_grad=0., weight=1., reduction=jnp.sum): """Represents an equality constraint, g(x) = 0. Args: @@ -177,7 +339,9 @@ def eq(fun,model_lagrangian='Standard', multiplier=0.0,penalty=1.,sq_grad=0., we """ def init_fn(*args, **kwargs): - return {'lambda': LagrangeMultiplier(multiplier+jnp.zeros_like(fun(*args, **kwargs)),penalty+jnp.zeros_like(fun(*args, **kwargs)),sq_grad+jnp.zeros_like(fun(*args, **kwargs)))} + out = fun(*args, **kwargs) + return {'lambda': _multiplier_like(out, multiplier, penalty, omega, eta, sq_grad)} + #return {'lambda': LagrangeMultiplier(multiplier+jnp.zeros_like(fun(*args, **kwargs)),penalty+jnp.zeros_like(fun(*args, **kwargs)),sq_grad+jnp.maximum(jnp.square(fun(*args, **kwargs)),1.e-4))} if model_lagrangian=='Standard': def loss_fn(params, *args, **kwargs): @@ -188,10 +352,10 @@ def loss_fn(params, *args, **kwargs): inf = fun(*args, **kwargs) return weight * reduction(-params['lambda'].value * inf + params['lambda'].penalty* inf ** 2 / 2+ params['lambda'].value**2 /(2.*params['lambda'].penalty)), inf - return Constraint(init_fn, loss_fn) + return BaseConstraint(init_fn, loss_fn) -def ineq(fun, model_lagrangian='Standard', multiplier=0.,penalty=1., sq_grad=0.,weight=1., reduction=jnp.sum): +def ineq(fun, model_lagrangian='Standard', multiplier=0.,penalty=1.,omega=1.0,eta=1.0, sq_grad=0.,weight=1., reduction=jnp.sum): """Represents an inequality constraint, h(x) >= 0, which uses a slack variable internally to convert it to an equality constraint. @@ -212,8 +376,8 @@ def ineq(fun, model_lagrangian='Standard', multiplier=0.,penalty=1., sq_grad=0., def init_fn(*args, **kwargs): out = fun(*args, **kwargs) - return {'lambda': LagrangeMultiplier(multiplier+jnp.zeros_like(fun(*args, **kwargs)),penalty+jnp.zeros_like(fun(*args, **kwargs)),sq_grad+jnp.zeros_like(fun(*args, **kwargs))), - 'slack': jax.nn.relu(out) ** 0.5} + return {'lambda': _multiplier_like(out, multiplier, penalty, omega, eta, sq_grad), + 'slack': jax.nn.relu(out) ** 0.5} if model_lagrangian=='Standard': def loss_fn(params, *args, **kwargs): @@ -224,28 +388,166 @@ def loss_fn(params, *args, **kwargs): inf = fun(*args, **kwargs) - params['slack'] ** 2 return weight * reduction(-params['lambda'].value * inf + params['lambda'].penalty * inf ** 2 / 2+ params['lambda'].value**2 /(2.*params['lambda'].penalty)), inf - return Constraint(init_fn, loss_fn) + return BaseConstraint(init_fn, loss_fn) def combine(*args): - """Combines multiple constraint tuples into a single constraint tuple. + """Combines constraints with selective named dependencies, mirroring losses.py. + + Each SelectiveConstraint specifies which named arguments it depends on. + Regular Constraints still work (they receive all arguments positionally). + + The returned combined constraint supports both: + - Old style: `combined.init(arg1, arg2); combined.loss(params, arg1, arg2)` + - New style: Set `combined.dependencies = {...}` and call `combined.init/loss()` + + Implementation optimizes for JIT by pre-wrapping constraint functions at combination time, + avoiding dynamic control flow inside the loss/init functions. + + Validation: + - All constraints must be Constraint or SelectiveConstraint objects + - Constraint outputs must be compatible (all scalars, all same shape, etc.) Args: - *args: A series of constraint (init_fn, loss_fn) tuples. + *args: A series of constraint (init_fn, loss_fn) tuples or SelectiveConstraint objects. Returns: - A single (init_fn, loss_fn) tuple that wraps the input constraints. + A combined Constraint with optional `.dependencies` and `.arg_names` attributes. + + Raises: + ValueError: If no constraints provided or constraint types are invalid + TypeError: If constraint objects are not Constraint or SelectiveConstraint """ - init_fns, loss_fns = zip(*args) + if not args: + raise ValueError("At least one constraint must be provided to combine()") + + # Separate init_fns and loss_fns, tracking SelectiveConstraints + constraints_list = [] + selective_map = {} # Maps position to SelectiveConstraint + all_arg_names = [] # ordered list of unique arg names for named dependency mode + + for i, arg in enumerate(args): + # Extract the underlying callable pair and record selective wrappers + if isinstance(arg, SelectiveConstraint): + c = arg.constraint + selective_map[i] = arg + # preserve order and uniqueness of arg names + for name in arg.arg_names: + if name not in all_arg_names: + all_arg_names.append(name) + else: + c = arg + # Accept tuple (init_fn, loss_fn) or objects with .init and .loss + if isinstance(c, tuple) and len(c) == 2: + init_fn, loss_fn = c + elif hasattr(c, 'init') and hasattr(c, 'loss'): + init_fn, loss_fn = c.init, c.loss + else: + raise TypeError( + f"Constraint {i} must be SelectiveConstraint, BaseConstraint-like, or (init_fn, loss_fn) tuple, got {type(arg)}" + ) + + constraints_list.append((init_fn, loss_fn)) + + init_fns, loss_fns = zip(*constraints_list) + + # Pre-wrap constraint functions to bake in filtering at combination time + # This avoids dynamic control flow (if/else) inside JIT-compiled functions + wrapped_init_fns = [] + wrapped_loss_fns = [] + + for i, (init_fn, loss_fn) in enumerate(zip(init_fns, loss_fns)): + if i in selective_map: + selective = selective_map[i] + + # Wrap init_fn to use named dependencies + def make_init_wrapper(s, f): + def wrapped_init(*a, **kw): + # Merge constructor kwargs with call-time kwargs, prefer call-time + merged_kw = {**s.kwargs, **kw} + if a or kw: + # If first arg looks like the flat dofs (array or tracer), map it + first = a[0] if a else None + is_object_like = hasattr(first, 'B') or hasattr(first, 'coils') or hasattr(first, 'dofs') or isinstance(first, dict) + looks_flat = (hasattr(first, 'ndim') or hasattr(first, 'shape') or hasattr(first, 'aval')) + if a and (looks_flat and not is_object_like): + # Try per-selective unravel first; if it fails, fall back to composite unravel + try: + pytrees = s.dofs_to_pytree(a[0]) + return f(*pytrees, **merged_kw) + except Exception: + # fall back to composite-level mapping if available + if hasattr(s, '_composite_index_map') and hasattr(combined, 'dofs_to_pytree'): + all_pytrees = combined.dofs_to_pytree(a[0]) + pytrees = tuple(all_pytrees[idx] for idx in s._composite_index_map) + return f(*pytrees, **merged_kw) + raise + if is_object_like: + return f(*a, **merged_kw) + # fall through to use stored dependencies + filtered_args = s._get_filtered_args() + return f(*filtered_args, **s.kwargs) + return wrapped_init + wrapped_init_fns.append(make_init_wrapper(selective, init_fn)) + + # Wrap loss_fn to use named dependencies + def make_loss_wrapper(s, f): + def wrapped_loss(p, *a, **kw): + merged_kw = {**s.kwargs, **kw} + if a or kw: + first = a[0] if a else None + is_object_like = hasattr(first, 'B') or hasattr(first, 'coils') or hasattr(first, 'dofs') or isinstance(first, dict) + looks_flat = (hasattr(first, 'ndim') or hasattr(first, 'shape') or hasattr(first, 'aval')) + if a and (looks_flat and not is_object_like): + try: + pytrees = s.dofs_to_pytree(a[0]) + return f(p, *pytrees, **merged_kw) + except Exception: + if hasattr(s, '_composite_index_map') and hasattr(combined, 'dofs_to_pytree'): + all_pytrees = combined.dofs_to_pytree(a[0]) + pytrees = tuple(all_pytrees[idx] for idx in s._composite_index_map) + return f(p, *pytrees, **merged_kw) + raise + if is_object_like: + return f(p, *a, **merged_kw) + # fall through to stored dependencies + filtered_args = s._get_filtered_args() + return f(p, *filtered_args, **s.kwargs) + return wrapped_loss + wrapped_loss_fns.append(make_loss_wrapper(selective, loss_fn)) + else: + # Regular constraints: pass-through wrappers that accept all args + def make_init_wrapper_all(f): + def wrapped_init(*args, **kwargs): + return f(*args, **kwargs) + return wrapped_init + wrapped_init_fns.append(make_init_wrapper_all(init_fn)) + + def make_loss_wrapper_all(f): + def wrapped_loss(p, *args, **kwargs): + return f(p, *args, **kwargs) + return wrapped_loss + wrapped_loss_fns.append(make_loss_wrapper_all(loss_fn)) + + # Now init_fn and loss_fn are clean - no control flow, just list comprehensions def init_fn(*args, **kwargs): - return tuple(fn(*args, **kwargs) for fn in init_fns) + """Initialize all constraints using pre-wrapped functions.""" + results = [fn(*args, **kwargs) for fn in wrapped_init_fns] + return tuple(results) def loss_fn(params, *args, **kwargs): - outs = [fn(p, *args, **kwargs) for p, fn in zip(params, loss_fns)] + """Compute total loss from all constraints using pre-wrapped functions.""" + outs = [fn(p, *args, **kwargs) for p, fn in zip(params, wrapped_loss_fns)] return sum(x[0] for x in outs), tuple(x[1] for x in outs) - return Constraint(init_fn, loss_fn) + combined = CompositeConstraint(init_fn, loss_fn, selective_map=selective_map, arg_names=all_arg_names) + + # Precompute mapping from composite arg_names -> indices for each selective + for sel in selective_map.values(): + sel._composite_index_map = tuple(combined.arg_names.index(n) for n in sel.arg_names) + + return combined @@ -271,6 +573,178 @@ def penalty_average(tree): return jnp.average(penalty[0]) +def apply_mu_tolerance_per_constraint(constraint_dict, grad_dict, constraint_info, constraint_info_prev=None, model_lagrangian='Standard', model_mu='Mu_Adaptative_1', beta=2.0, mu_max=1.e4, alpha=0.99, gamma=1.e-2, epsilon=1.e-8, eta_tol=1.e-4, omega_tol=1.e-6, decrease_tol=0.75): + """Apply Mu update rule for vectorized constraints. + + Supports three strategies: + - Mu_Tolerance: All elements updated together based on norm, uses average penalty + - Mu_Adaptative_1: Element-wise updates using eta parameter + - Mu_Adaptative_2: Element-wise updates using decrease tolerance criterion + + Args: + constraint_dict: dict with 'lambda' and optionally 'slack' keys containing LagrangeMultiplier objects + grad_dict: dict with gradients matching the constraint_dict structure + constraint_info: current constraint violation information (array or scalar) + constraint_info_prev: previous constraint violation info (needed for Mu_Adaptative_2) + model_lagrangian: 'Standard' or 'Squared' lagrangian formulation + model_mu: 'Mu_Tolerance' (global norm, avg penalty), 'Mu_Adaptative_1' (element-wise eta), or 'Mu_Adaptative_2' (element-wise decrease) + decrease_tol: tolerance for decrease criterion in Mu_Adaptative_2 (default 0.75 = 25% decrease) + """ + pred = lambda x: isinstance(x, LagrangeMultiplier) + + # Extract eta and omega from the first LagrangeMultiplier + first_key = list(constraint_dict.keys())[0] + eta_val = constraint_dict[first_key].eta + omega_val = constraint_dict[first_key].omega + + # Extract penalty values + constraint_penalty = jax.tree_util.tree_map(lambda x: x.penalty, constraint_dict, is_leaf=pred) + penalty_val = constraint_penalty[first_key] + + # Get constraint absolute values (element-wise) + constraint_abs = jnp.abs(constraint_info) + + + # Element-wise strategies (Mu_Adaptative_1 or Mu_Adaptative_2) + # Compute updated eta and omega (element-wise) + eta_updated = jnp.maximum(eta_val / jnp.power(penalty_val, 0.1), eta_tol) + omega_updated = jnp.maximum(omega_val / penalty_val, omega_tol) + + eta_updated_false = jnp.maximum(1. / jnp.power(penalty_val, 0.1), eta_tol) + omega_updated_false = jnp.maximum(1.0 / penalty_val, omega_tol) + + # Determine if constraints are satisfied based on model_mu strategy + if model_mu == 'Mu_Adaptative_1': + # Strategy 1: Use eta parameter - constraint satisfied if abs(constraint) < eta + is_satisfied = constraint_abs < eta_val + + elif model_mu == 'Mu_Adaptative_2': + # Strategy 2: Use decrease tolerance - constraint satisfied if abs(constraint) <= decrease_tol * abs(constraint_prev) + if constraint_info_prev is None: + # If no previous constraint, fall back to eta-based check + is_satisfied = constraint_abs < eta_val + else: + constraint_abs_prev = jnp.abs(constraint_info_prev) + # Check if constraint decreased by specified tolerance factor + is_satisfied = constraint_abs <= decrease_tol * constraint_abs_prev + else: + # Default to Mu_Adaptative_1 + is_satisfied = constraint_abs < eta_val + + if model_lagrangian == 'Standard': + # For satisfied constraints: update lambda value, set penalty change to 0 + # For unsatisfied constraints: set lambda to 0, increase penalty + updated_dict = jax.tree_util.tree_map( + lambda x, y: LagrangeMultiplier( + jnp.where(is_satisfied, x.penalty * y.value, 0.0 * x.value), # lambda update + jnp.where(is_satisfied, 0.0 * x.penalty, jnp.minimum(beta * x.penalty, mu_max) - x.penalty), # penalty change + jnp.where(is_satisfied, -x.omega+omega_updated, -x.omega+omega_updated_false), # omega + jnp.where(is_satisfied, -x.eta+eta_updated, -x.eta+eta_updated_false), # eta + 0.0 * x.sq_grad + ), + constraint_dict, grad_dict, is_leaf=pred + ) + + elif model_lagrangian == 'Squared': + # For Squared lagrangian: lambda update differs + updated_dict = jax.tree_util.tree_map( + lambda x, y: LagrangeMultiplier( + jnp.where(is_satisfied, x.penalty * (y.value - x.value / x.penalty), 0.0 * x.value), # lambda update + jnp.where(is_satisfied, 0.0 * x.penalty, jnp.minimum(beta * x.penalty, mu_max) - x.penalty), # penalty change + jnp.where(is_satisfied, -x.omega+omega_updated, -x.omega+omega_updated_false), # omega + jnp.where(is_satisfied, -x.eta+eta_updated, -x.eta+eta_updated_false), # eta + 0.0 * x.sq_grad + ), + constraint_dict, grad_dict, is_leaf=pred + ) + + return updated_dict + + +def apply_mu_tolerance_all_constraints(constraint_dicts_map, grad_dicts_map=None, constraint_infos_map=None, model_lagrangian='Standard', beta=2.0, mu_max=1.e4, alpha=0.99, gamma=1.e-2, epsilon=1.e-8, eta_tol=1.e-4, omega_tol=1.e-6): + """Apply Mu_Tolerance across all constraints (JAX-compatible, no Python loops or dict indexing). + + This mirrors the per-constraint updater but computes a single global + decision using all constraint infos, then applies the same update to + every constraint's Lagrange multipliers. + Args: + constraint_dicts_map: pytree mapping keys -> per-constraint lagrange dicts + grad_dicts_map: optional pytree of matching shapes with gradient/info dicts + constraint_infos_map: optional pytree mapping keys -> constraint info pytrees + Returns: + pytree with updated LagrangeMultiplier dicts matching `constraint_dicts_map`. + """ + pred = lambda x: isinstance(x, LagrangeMultiplier) + + # Build global flat vector of all constraint infos + if constraint_infos_map is None: + global_norm = jnp.array(0.0) + else: + flat_map = jax.tree_util.tree_map(lambda info: jax.flatten_util.ravel_pytree(info)[0], constraint_infos_map) + leaves = jax.tree_util.tree_leaves(flat_map) + global_norm = jnp.linalg.norm(jnp.concatenate(leaves) if leaves else jnp.array(0.0)) + + # Extract eta/omega/penalty from all LagrangeMultipliers directly - no nested tree_map + # This matches the pattern in penalty_average function which works inside JIT + eta_pytree = jax.tree_util.tree_map(lambda x: x.eta, constraint_dicts_map, is_leaf=pred) + omega_pytree = jax.tree_util.tree_map(lambda x: x.omega, constraint_dicts_map, is_leaf=pred) + penalty_pytree = jax.tree_util.tree_map(lambda x: x.penalty, constraint_dicts_map, is_leaf=pred) + + # Flatten to get scalar values + eta_flat = jax.flatten_util.ravel_pytree(eta_pytree)[0] + omega_flat = jax.flatten_util.ravel_pytree(omega_pytree)[0] + penalty_flat = jax.flatten_util.ravel_pytree(penalty_pytree)[0] + + global_eta = jnp.mean(eta_flat) if eta_flat.size > 0 else eta_tol + global_omega = jnp.mean(omega_flat) if omega_flat.size > 0 else omega_tol + mu_average = jnp.mean(penalty_flat) if penalty_flat.size > 0 else 1.0 + + is_satisfied = global_norm < global_eta + + # Compute updated eta and omega based on satisfaction + eta_updated = jnp.maximum(global_eta / jnp.power(mu_average, 0.1), eta_tol) + omega_updated = jnp.maximum(global_omega / mu_average, omega_tol) + + eta_updated_false = jnp.maximum(1. / jnp.power(mu_average, 0.1), eta_tol) + omega_updated_false = jnp.maximum(1.0 / mu_average, omega_tol) + + # Prepare grad_map fallback (zeroed) if not provided + if grad_dicts_map is None: + def _zero_leaf(leaf): + return LagrangeMultiplier(jnp.zeros_like(leaf.value), jnp.zeros_like(leaf.penalty), jnp.zeros_like(leaf.omega), jnp.zeros_like(leaf.eta), jnp.zeros_like(leaf.sq_grad)) + grad_map_full = jax.tree_util.tree_map(lambda c: jax.tree_util.tree_map(_zero_leaf, c, is_leaf=pred), constraint_dicts_map) + else: + grad_map_full = grad_dicts_map + + def _update_single(cdict, gdict): + if model_lagrangian == 'Standard': + return jax.tree_util.tree_map( + lambda x, y: LagrangeMultiplier( + jnp.where(is_satisfied, x.penalty * y.value, 0.0 * x.value), + jnp.where(is_satisfied, 0.0 * x.penalty, jnp.minimum(beta * x.penalty, mu_max) - x.penalty), + jnp.where(is_satisfied, -x.omega + omega_updated, -x.omega + omega_updated_false), + jnp.where(is_satisfied, -x.eta + eta_updated, -x.eta + eta_updated_false), + 0.0 * x.sq_grad + ), + cdict, gdict, is_leaf=pred + ) + else: + return jax.tree_util.tree_map( + lambda x, y: LagrangeMultiplier( + jnp.where(is_satisfied, x.penalty * (y.value - x.value / x.penalty), 0.0 * x.value), + jnp.where(is_satisfied, 0.0 * x.penalty, jnp.minimum(beta * x.penalty, mu_max) - x.penalty), + jnp.where(is_satisfied, -x.omega + omega_updated, -x.omega + omega_updated_false), + jnp.where(is_satisfied, -x.eta + eta_updated, -x.eta + eta_updated_false), + 0.0 * x.sq_grad + ), + cdict, gdict, is_leaf=pred + ) + + # Map over constraint dicts, treating dicts as leaves (not LagrangeMultipliers) + updated_map = jax.tree_util.tree_map(_update_single, constraint_dicts_map, grad_map_full, is_leaf=lambda x: isinstance(x, dict)) + return updated_map + + @@ -285,10 +759,10 @@ class ALM(NamedTuple): #This can use optax gradient descent optimizers with different mu updating methods def ALM_model_optax(optimizer: optax.GradientTransformation, #an optimizer from OPTAX - constraints: Constraint, #List of constraints + constraints: BaseConstraint, #List of constraints loss= lambda x: 0., #function which represents the loss (Callable, default 0.) model_lagrangian='Standard' , #Model to use for updating lagrange multipliers - model_mu='Constant' , #Model to use for updating lagrange multipliers + model_mu='Mu_Tolerance' , #Model to use for updating lagrange multipliers beta=2.0, mu_max=1.e4, alpha=0.99, @@ -299,87 +773,48 @@ def ALM_model_optax(optimizer: optax.GradientTransformation, #an optimizer from **kargs, #Extra key arguments for loss ): - - if model_mu=='Mu_Tolerance_LBFGS': - @jax.jit - def init_fn(params,**kargs): - main_params,lagrange_params=params - main_state = optimizer.init(main_params) - lag_state=lagrange_update(model_lagrangian=model_lagrangian).init(lagrange_params) - opt_state=main_state,lag_state - value,grad=jax.value_and_grad(lagrangian,has_aux=True,argnums=(0,1))(main_params,lagrange_params,**kargs) - return opt_state,grad,value[0],value[1] - else: - @jax.jit - def init_fn(params,**kargs): - main_params,lagrange_params=params - main_state = optimizer.init(main_params) - lag_state=lagrange_update(model_lagrangian=model_lagrangian).init(lagrange_params) - opt_state=main_state,lag_state - grad,info=jax.grad(lagrangian,has_aux=True,argnums=(0,1))(main_params,lagrange_params,**kargs) - return opt_state,grad,info + @jax.jit + def init_fn(params,**kargs): + main_params,lagrange_params=params + main_state = optimizer.init(main_params) + lag_state=None + grad,info=jax.grad(lagrangian,has_aux=True,argnums=(0,1))(main_params,lagrange_params,**kargs) + return (main_state,lag_state),grad,info # Define the Augmented lagrangian if model_lagrangian=='Standard': def lagrangian(main_params,lagrange_params,**kargs): - main_loss = jnp.linalg.norm(loss(main_params,**kargs)) #The norm here is to ensure we have a scalr from the loss which should be a vector + main_loss = jnp.linalg.norm(loss(main_params,**kargs)) mdmm_loss, inf = constraints.loss(lagrange_params, main_params) return main_loss+mdmm_loss, (main_loss,main_loss+mdmm_loss, inf) - # Augmented Lagrangian - def lagrangian_lbfgs(main_params,lagrange_params,**kargs): - main_loss = jnp.linalg.norm(loss(main_params,**kargs)) - mdmm_loss, _ = constraints.loss(lagrange_params, main_params) - return main_loss+mdmm_loss - elif model_lagrangian=='Squared': def lagrangian(main_params,lagrange_params,**kargs): main_loss = jnp.square(jnp.linalg.norm(loss(main_params,**kargs))) - #Here we take the square because the term appearing in this Lagrangian mdmm_loss, inf = constraints.loss(lagrange_params, main_params) return main_loss+mdmm_loss, (main_loss,main_loss+mdmm_loss, inf) - # Augmented Lagrangian - def lagrangian_lbfgs(main_params,lagrange_params,**kargs): - #Here we take the square because the term appearing in this Lagrangian - main_loss = jnp.square(jnp.linalg.norm(loss(main_params,**kargs))) - mdmm_loss, _ = constraints.loss(lagrange_params, main_params) - return main_loss+mdmm_loss - - if model_mu=='Mu_Conditional': - # Do the optimization step - @partial(jit, static_argnums=(6,7,8,9,10,11,12,13)) - def update_fn(params, opt_state,grad,info,eta,omega,model_lagrangian=model_lagrangian,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol,**kargs): + if model_mu=='Mu_Adaptative_2': + @partial(jit, static_argnums=(4,5,6,7,8,9,10)) + def update_fn(params, opt_state,grad,info,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol,**kargs): main_state,lag_state=opt_state main_params,lagrange_params=params - main_updates, main_state = optimizer.update(grad[0], main_state) - main_params = optax.apply_updates(main_params, main_updates) - params=main_params,lagrange_params - grad,info = jax.grad(lagrangian,has_aux=True,argnums=(0,1))(main_params,lagrange_params,**kargs) - true_func=partial(lagrange_update(model_lagrangian=model_lagrangian).update,model_mu='Mu_Conditional_True',beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) - false_func=partial(lagrange_update(model_lagrangian=model_lagrangian).update,model_mu='Mu_Conditional_False',beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) - lag_updates, lag_state = jax.lax.cond(norm_constraints(info[2]) omega + return jnp.linalg.norm(grad[0])> omega_min def minimization_loop(state): params,main_state,grad,info=state main_params,lagrange_params=params - #jax.debug.print('Loop omega: {omega}', omega=omega) - #jax.debug.print('Loop grad: {grad}', grad=jnp.linalg.norm(grad[0])) main_updates, main_state = optimizer.update(grad[0], main_state) main_params = optax.apply_updates(main_params, main_updates) params=main_params,lagrange_params @@ -389,74 +824,98 @@ def minimization_loop(state): params,main_state,grad,info=jax.lax.while_loop(condition,minimization_loop,state) main_params,lagrange_params=params - true_func=partial(lagrange_update(model_lagrangian=model_lagrangian).update,model_mu='Mu_Tolerance_True',beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) - false_func=partial(lagrange_update(model_lagrangian=model_lagrangian).update,model='Mu_Tolerance_False',beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) - lag_updates, lag_state = jax.lax.cond(norm_constraints(info[2]) omega + _,_,grad,_=state + return jnp.linalg.norm(grad[0])> omega_min def minimization_loop(state): - params,main_state,grad,value,info=state + params,main_state,grad,info=state main_params,lagrange_params=params - #jax.debug.print('Loop omega: {omega}', omega=omega) - #jax.debug.print('Loop grad: {grad}', grad=jnp.linalg.norm(grad[0])) - main_updates, main_state = optimizer.update(grad[0], main_state,params=main_params,value=value,grad=grad[0],value_fn=lagrangian_lbfgs,lagrange_params=lagrange_params) + main_updates, main_state = optimizer.update(grad[0], main_state) main_params = optax.apply_updates(main_params, main_updates) params=main_params,lagrange_params - value,grad = jax.value_and_grad(lagrangian,has_aux=True,argnums=(0,1))(main_params,lagrange_params,**kargs) - #Here info is in value[1] - state=params,main_state,grad,value[0],value[1] + grad,info = jax.grad(lagrangian,has_aux=True,argnums=(0,1))(main_params,lagrange_params,**kargs) + state=params,main_state,grad,info return state - params,main_state,grad,value,info=jax.lax.while_loop(condition,minimization_loop,state) + params,main_state,grad,info=jax.lax.while_loop(condition,minimization_loop,state) main_params,lagrange_params=params - true_func=partial(lagrange_update(model_lagrangian=model_lagrangian).update,model_mu='Mu_Tolerance_True',beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) - false_func=partial(lagrange_update(model_lagrangian=model_lagrangian).update,model_mu='Mu_Tolerance_False',beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) - lag_updates, lag_state = jax.lax.cond(norm_constraints(info[2]) 2, "n_segments must be greater than 2" assert isinstance(nfp, int), "nfp must be a positive integer" assert nfp > 0, "nfp must be a positive integer" assert isinstance(stellsym, bool), "stellsym must be a boolean" - + self._initialize_state( + dofs, + n_segments, + nfp, + stellsym, + self._normalize_scaling_type(scaling_type), + scaling_factor, + scale_fixed, + ) + + def _initialize_state(self, dofs, n_segments, nfp, stellsym, scaling_type, scaling_factor, scale_fixed, order=None): self._dofs = dofs self._n_segments = n_segments self._nfp = nfp self._stellsym = stellsym - self._order = dofs.shape[2]//2 - self._curves = apply_symmetries_to_curves(self.dofs, self.nfp, self.stellsym) - self.quadpoints = jnp.linspace(0, 1, self.n_segments, endpoint=False) - self._set_gamma() - self.n_base_curves=dofs.shape[0] - - def __str__(self): - return f"nfp stellsym order\n{self.nfp} {self.stellsym} {self.order}\n"\ - + f"Degrees of freedom\n{repr(self.dofs.tolist())}\n" - - def __repr__(self): - return f"nfp stellsym order\n{self.nfp} {self.stellsym} {self.order}\n"\ - + f"Degrees of freedom\n{repr(self.dofs.tolist())}\n" + self._scaling_type = scaling_type + self._scaling_factor = scaling_factor + self._scale_fixed = scale_fixed + self._scaling = None + self._order = dofs.shape[2] // 2 if hasattr(dofs, "shape") else order + self.quadpoints = jnp.linspace(0, 1, self._n_segments, endpoint=False) + self._curves = None + self._gamma = None + self._gamma_dash = None + self._gamma_dashdash = None + self._length = None + self._curvature = None - def _tree_flatten(self): - children = (self._dofs,) # arrays / dynamic values - aux_data = {"n_segments": self._n_segments, "nfp": self._nfp, "stellsym": self._stellsym} # static values - return (children, aux_data) + @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 _tree_unflatten(cls, aux_data, children): - return cls(*children, **aux_data) - - partial(jit, static_argnames=['self']) - def _set_gamma(self): - def fori_createdata(order_index: int, data: jnp.ndarray) -> jnp.ndarray: - return data[0] + jnp.einsum("ij,k->ikj", self._curves[:, :, 2 * order_index - 1], jnp.sin(2 * jnp.pi * order_index * self.quadpoints)) + jnp.einsum("ij,k->ikj", self._curves[:, :, 2 * order_index], jnp.cos(2 * jnp.pi * order_index * self.quadpoints)), \ - data[1] + jnp.einsum("ij,k->ikj", self._curves[:, :, 2 * order_index - 1], 2*jnp.pi *order_index *jnp.cos(2 * jnp.pi * order_index * self.quadpoints)) + jnp.einsum("ij,k->ikj", self._curves[:, :, 2 * order_index], -2*jnp.pi *order_index *jnp.sin(2 * jnp.pi * order_index * self.quadpoints)), \ - data[2] + jnp.einsum("ij,k->ikj", self._curves[:, :, 2 * order_index - 1], -4*jnp.pi**2*order_index**2*jnp.sin(2 * jnp.pi * order_index * self.quadpoints)) + jnp.einsum("ij,k->ikj", self._curves[:, :, 2 * order_index], -4*jnp.pi**2*order_index**2*jnp.cos(2 * jnp.pi * order_index * self.quadpoints)) - gamma = jnp.einsum("ij,k->ikj", self._curves[:, :, 0], jnp.ones(self.n_segments)) - gamma_dash = jnp.zeros((jnp.size(self._curves, 0), self.n_segments, 3)) - gamma_dashdash = jnp.zeros((jnp.size(self._curves, 0), self.n_segments, 3)) - gamma, gamma_dash, gamma_dashdash = fori_loop(1, self._order+1, fori_createdata, (gamma, gamma_dash, gamma_dashdash)) - length = jnp.array([jnp.mean(jnp.linalg.norm(d1gamma, axis=1)) for d1gamma in gamma_dash]) - curvature = vmap(compute_curvature)(gamma_dash, gamma_dashdash) - self._gamma = gamma - self._gamma_dash = gamma_dash - self._gamma_dashdash = gamma_dashdash - self._curvature = curvature - self._length = length + @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): + self._curves = None + self._gamma = None + self._gamma_dash = None + self._gamma_dashdash = None + self._curvature = None + self._length = None + + # dofs property and setter @property def dofs(self): - return 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): - assert isinstance(new_dofs, jnp.ndarray) - assert new_dofs.ndim == 3 - assert jnp.size(new_dofs, 1) == 3 - assert jnp.size(new_dofs, 2) % 2 == 1 - self._dofs = new_dofs - self._order = jnp.size(new_dofs, 2)//2 - self._curves = apply_symmetries_to_curves(self.dofs, self.nfp, self.stellsym) - self._set_gamma() + self.reset_cache() + self._dofs = new_dofs / self.scaling[None, None, :] + self._order = self._dofs.shape[2] // 2 - @property - def curves(self): - return self._curves - - @property - def order(self): - return self._order - - @order.setter - def order(self, new_order): - assert isinstance(new_order, int) - assert new_order > 0 - self._dofs = jnp.pad(self.dofs, ((0, 0), (0, 0), (0, 2*(new_order-self._order)))) if new_order > self._order else self.dofs[:, :, :2*(new_order)+1] - self._order = new_order - self._curves = apply_symmetries_to_curves(self.dofs, self.nfp, self.stellsym) - self._set_gamma() - + # n_segments property and setter @property def n_segments(self): return self._n_segments @n_segments.setter def n_segments(self, new_n_segments): - assert isinstance(new_n_segments, int) - assert new_n_segments > 2 + self.reset_cache() self._n_segments = new_n_segments self.quadpoints = jnp.linspace(0, 1, self._n_segments, endpoint=False) - self._set_gamma() - + + # nfp property and setter @property def nfp(self): return self._nfp @nfp.setter def nfp(self, new_nfp): - assert isinstance(new_nfp, int) - assert new_nfp > 0 + self.reset_cache() self._nfp = new_nfp - self._curves = apply_symmetries_to_curves(self.dofs, self.nfp, self.stellsym) - self._set_gamma() - + + # stellsym property and setter @property def stellsym(self): return self._stellsym @stellsym.setter def stellsym(self, new_stellsym): - assert isinstance(new_stellsym, bool) + self.reset_cache() self._stellsym = new_stellsym - self._curves = apply_symmetries_to_curves(self.dofs, self.nfp, self.stellsym) - self._set_gamma() + # 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: + scaling = self._compute_mode_scaling( + self.order, self.scaling_type, self.scaling_factor, self.scale_fixed + ) + if not isinstance(scaling, jax.core.Tracer): + self._scaling = scaling + return scaling + return self._scaling + + # order property and setter + @property + def order(self): + if hasattr(self._dofs, "shape"): + return self._dofs.shape[2] // 2 + return self._order + + @order.setter + def order(self, new_order): + self.reset_cache() + # Get unscaled dofs, resize, then store unscaled + 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 + self._order = new_order + + # n_base_curves property + @property + def n_base_curves(self): + return self._dofs.shape[0] + + # curves property + @property + def curves(self): + if self._curves is None: + # 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 + @jit + def _compute_gamma(self): + def create_data(order: int) -> jnp.ndarray: + return jnp.einsum("ij,k->ikj", self.curves[:, :, 2 * order - 1], jnp.sin(2 * jnp.pi * order * self.quadpoints)) \ + + jnp.einsum("ij,k->ikj", self.curves[:, :, 2 * order], jnp.cos(2 * jnp.pi * order * self.quadpoints)) + gamma_0 = jnp.einsum("ij,k->ikj", self.curves[:, :, 0], jnp.ones(self.n_segments)) + gamma_n = vmap(create_data)(jnp.arange(1, self.order+1)) + return gamma_0 + jnp.sum(gamma_n, axis=0) + + # TODO change gamma from a property to a method + # gamma property @property def gamma(self): + if self._gamma is None: + self._gamma = self._compute_gamma() return self._gamma - - @gamma.setter - def gamma(self, new_gamma): - self._gamma = new_gamma + # _compute_gamma_dash method + @jit + def _compute_gamma_dash(self): + def create_data(order: int) -> jnp.ndarray: + return jnp.einsum("ij,k->ikj", self.curves[:, :, 2 * order - 1], 2*jnp.pi * order * jnp.cos(2 * jnp.pi * order * self.quadpoints)) \ + + jnp.einsum("ij,k->ikj", self.curves[:, :, 2 * order], -2 * jnp.pi * order * jnp.sin(2 * jnp.pi * order * self.quadpoints)) + gamma_dash_n = vmap(create_data)(jnp.arange(1, self.order+1)) + return jnp.sum(gamma_dash_n, axis=0) + + # 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_dash.setter - def gamma_dash(self, new_gamma_dash): - self._gamma_dash = new_gamma_dash - + # _compute_gamma_dashdash method + @jit + def _compute_gamma_dashdash(self): + def create_data(order: int) -> jnp.ndarray: + return jnp.einsum("ij,k->ikj", self.curves[:, :, 2 * order - 1], -4*jnp.pi**2 * order**2 * jnp.sin(2 * jnp.pi * order * self.quadpoints)) \ + + jnp.einsum("ij,k->ikj", self.curves[:, :, 2 * order], -4*jnp.pi**2 * order**2 * jnp.cos(2 * jnp.pi * order * self.quadpoints)) + gamma_dashdash_n = vmap(create_data)(jnp.arange(1, self.order+1)) + return jnp.sum(gamma_dashdash_n, axis=0) - + # 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 - @gamma_dashdash.setter - def gamma_dashdash(self, new_gamma_dashdash): - self._gamma_dashdash = new_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 + # compute_curvature static method + @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 + + # curvature property @property def curvature(self): - return self._curvature + return vmap(self.compute_curvature)(self.gamma_dash, self.gamma_dashdash) + + # copy method + def copy(self): + curves = object.__new__(Curves) + curves._initialize_state( + self._dofs.copy(), + self._n_segments, + self._nfp, + self._stellsym, + self._scaling_type, + self._scaling_factor, + self._scale_fixed, + order=self.order, + ) + return curves + + # magic methods + def __str__(self): + return f"nfp stellsym order\n{self.nfp} {self.stellsym} {self.order}\n"\ + + f"Degrees of freedom\n{repr(self.dofs.tolist())}\n" + + def __repr__(self): + return f"nfp stellsym order\n{self.nfp} {self.stellsym} {self.order}\n"\ + + f"Degrees of freedom\n{repr(self.dofs.tolist())}\n" def __len__(self): - return jnp.size(self.curves, 0) + return self.curves.shape[0] def __getitem__(self, key): if isinstance(key, int): @@ -308,92 +450,306 @@ def wrap(data): pointData = {**pointData, **extra_data} polyLinesToVTK(str(filename), np.array(x), np.array(y), np.array(z), pointsPerLine=np.array(ppl), pointData=pointData) -class Curves_from_simsopt(Curves): - # This assumes curves have all nfp and stellsym symmetries - def __init__(self, simsopt_curves, nfp=1, stellsym=True): + @classmethod + 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 bs = load(simsopt_curves) simsopt_coils = bs.coils simsopt_curves = [c.curve for c in simsopt_coils] simsopt_curves = simsopt_curves[0:int(len(simsopt_curves)/nfp/(1+stellsym))] - dofs = jnp.reshape(jnp.array( - [curve.x for curve in simsopt_curves] + dofs = jnp.reshape(jnp.asarray( + [jnp.asarray(curve.x, dtype=float) for curve in simsopt_curves], + dtype=float, ), (len(simsopt_curves), 3, 2*simsopt_curves[0].order+1)) n_segments = len(simsopt_curves[0].quadpoints) - super().__init__(dofs, n_segments, nfp, stellsym) + return cls(dofs, n_segments, nfp, stellsym, scaling_type, scaling_factor, scale_fixed) + + def _tree_flatten(self): + dofs = self.dofs if hasattr(self._dofs, "shape") else self._dofs + children = (dofs,) # arrays / dynamic values + aux_data = {"n_segments": self._n_segments, + "nfp": self._nfp, + "stellsym": self._stellsym, + "scaling_type": self._scaling_type, + "scaling_factor": self._scaling_factor, + "scale_fixed": self._scale_fixed, + "order": self.order} # static values + return (children, aux_data) + + @classmethod + def _tree_unflatten(cls, aux_data, children): + dofs, = children + if hasattr(dofs, "shape"): + scaling = cls._compute_mode_scaling( + aux_data["order"], + aux_data["scaling_type"], + aux_data["scaling_factor"], + aux_data["scale_fixed"], + ) + dofs = dofs / scaling[None, None, :] + obj = object.__new__(cls) + obj._initialize_state( + dofs, + aux_data["n_segments"], + aux_data["nfp"], + aux_data["stellsym"], + aux_data["scaling_type"], + aux_data["scaling_factor"], + aux_data["scale_fixed"], + order=aux_data["order"], + ) + return obj tree_util.register_pytree_node(Curves, Curves._tree_flatten, Curves._tree_unflatten) -class Coils(Curves): - def __init__(self, curves: Curves, currents: jnp.ndarray): - assert isinstance(curves, Curves) - currents = jnp.array(currents) - assert jnp.size(currents) == jnp.size(curves.dofs, 0) - super().__init__(curves.dofs, curves.n_segments, curves.nfp, curves.stellsym) - self._currents_scale = jnp.mean(jnp.abs(currents)) - self._dofs_currents = currents/self._currents_scale - self._currents = apply_symmetries_to_currents(self._dofs_currents*self._currents_scale, self.nfp, self.stellsym) - def __str__(self): - return f"nfp stellsym order\n{self.nfp} {self.stellsym} {self.order}\n"\ - + f"Degrees of freedom\n{repr(self.dofs.tolist())}\n" \ - + f"Currents degrees of freedom\n{repr(self.dofs_currents.tolist())}\n" \ - + f"Currents scaling factor\n{self.currents_scale}\n" - - def __repr__(self): - return f"nfp stellsym order\n{self.nfp} {self.stellsym} {self.order}\n"\ - + f"Degrees of freedom\n{repr(self.dofs.tolist())}\n" \ - + f"Currents degrees of freedom\n{repr(self.dofs_currents.tolist())}\n" \ - + f"Currents scaling factor\n{self.currents_scale}\n" +def _initialize_currents_scale(currents, currents_scale): + """Return a fixed current scale for normalized current dofs.""" + currents = jnp.atleast_1d(jnp.asarray(currents)) + if currents_scale is None: + return jnp.mean(jnp.abs(currents)) + return currents_scale + +def _normalize_base_currents(currents, curves): + """Return base currents as a 1D array matching the number of base curves.""" + currents = jnp.atleast_1d(jnp.asarray(currents)) + if hasattr(curves._dofs, "shape"): + n_base_curves = curves._dofs.shape[0] + if currents.shape[0] == 1 and n_base_curves != 1: + currents = jnp.full((n_base_curves,), currents[0]) + return currents + +def _currents_as_array(currents): + if isinstance(currents, bool): + return None + if isinstance(currents, (list, tuple)) or hasattr(currents, "shape") or jnp.isscalar(currents): + return jnp.atleast_1d(jnp.asarray(currents)) + return None + + +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 + + Attributes: + curves (Curves): Curves object storing the coil geometry + dofs_currents_raw (jnp.ndarray - shape (n_base_curves,)): Non-normalized currents of the base curves + currents_scale (float): Normalization factor for the currents + dofs_currents (jnp.ndarray - shape (n_base_curves,)): Normalized currents of the base curves + currents (jnp.ndarray - shape (n_base_curves * nfp * (1 + stellsym),)): Currents obtained by applying symmetries to the base currents + dofs_curves (jnp.ndarray - shape (n_base_curves, 3, 2*order+1)): Degrees of freedom of the curves + 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, 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" + + currents_array = _currents_as_array(currents) + if currents_array is not None: + currents_scale = _initialize_currents_scale(currents_array, currents_scale) + self._initialize_state(curves, currents, currents_scale) + + def _initialize_state(self, curves, currents_raw, currents_scale): + self.curves = curves + currents_array = _currents_as_array(currents_raw) + if currents_array is not None: + currents_raw = currents_array + currents_raw = _normalize_base_currents(currents_raw, curves) + self._dofs_currents_raw = currents_raw + self._currents_scale = currents_scale + self._dofs_currents = None if hasattr(currents_raw, "shape") else currents_raw + self._currents = None + + # reset_cache method + def reset_cache(self): + self._dofs_currents = None + self._currents = None + + # dofs_curves property and setter @property def dofs_curves(self): - return self._dofs + return self.curves.dofs @dofs_curves.setter def dofs_curves(self, new_dofs_curves): - self.dofs = new_dofs_curves + self.curves.dofs = new_dofs_curves + + # 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): + self.reset_cache() + self._dofs_currents_raw = jnp.atleast_1d(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): + # Sentinel leaf during PyTree traversal: pass through, don't scale. + if self._dofs_currents_raw is None or isinstance(self._dofs_currents_raw, bool): + return self._dofs_currents_raw + if self._dofs_currents is None: + dofs_currents = self.dofs_currents_raw / self.currents_scale + if not isinstance(dofs_currents, jax.core.Tracer): + self._dofs_currents = dofs_currents + return dofs_currents return self._dofs_currents @dofs_currents.setter def dofs_currents(self, new_dofs_currents): - self._dofs_currents = new_dofs_currents - self._currents = apply_symmetries_to_currents(self._dofs_currents*self.currents_scale, self.nfp, self.stellsym) - + self.dofs_currents_raw = new_dofs_currents * self.currents_scale + + # dofs property and setter @property - def currents_scale(self): - return self._currents_scale + def dofs(self): + return jnp.hstack([self.dofs_curves.ravel(), self.dofs_currents]) - @currents_scale.setter - def currents_scale(self, new_currents_scale): - self._currents_scale = new_currents_scale - self._currents = apply_symmetries_to_currents(self.dofs_currents*new_currents_scale, self.nfp, self.stellsym) + @dofs.setter + def dofs(self, new_dofs): + n_curve_dofs = jnp.size(self.dofs_curves) + self.dofs_curves = jnp.reshape(new_dofs[:n_curve_dofs], self.dofs_curves.shape) + self.dofs_currents = new_dofs[n_curve_dofs:] + # TODO: remove x property. This is a placeholder for compatibility with the examples that need to be updated. + # x property and setter @property def x(self): - dofs_curves = jnp.ravel(self.dofs_curves) - dofs_currents = jnp.ravel(self.dofs_currents) - return jnp.concatenate((dofs_curves, dofs_currents)) + return self.dofs @x.setter def x(self, new_dofs): - old_dofs_curves = jnp.ravel(self.dofs) - old_dofs_currents = jnp.ravel(self.dofs_currents) - new_dofs_curves = new_dofs[:old_dofs_curves.shape[0]] - new_dofs_currents = new_dofs[old_dofs_curves.shape[0]:] - self.dofs_curves = jnp.reshape(new_dofs_curves, (self.dofs_curves.shape)) - self.dofs_currents = new_dofs_currents - + self.dofs = new_dofs + + # 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 + # gamma property + @property + def gamma(self): + return self.curves.gamma + + # gamma_dash property + @property + def gamma_dash(self): + return self.curves.gamma_dash + + # gamma_dashdash property + @property + def gamma_dashdash(self): + return self.curves.gamma_dashdash + + # length property + @property + def length(self): + return self.curves.length + + # curvature property + @property + def curvature(self): + return self.curves.curvature + + # nfp property + @property + def nfp(self): + return self.curves.nfp + + # stellsym property + @property + def stellsym(self): + return self.curves.stellsym + + # order property + @property + def order(self): + return self.curves.order + + # n_segments property and setter + @property + def n_segments(self): + return self.curves.n_segments + + @n_segments.setter + def n_segments(self, new_n_segments): + self.curves.n_segments = new_n_segments + + # copy method + def copy(self): + 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 = self._currents + + return coils + + # magic methods + def __str__(self): + return f"nfp stellsym order\n{self.nfp} {self.stellsym} {self.order}\n"\ + + f"Degrees of freedom\n{repr(self.dofs.tolist())}\n" \ + + f"Currents degrees of freedom\n{repr(self.dofs_currents.tolist())}\n" \ + + f"Currents scaling factor\n{self.currents_scale}\n" + + def __repr__(self): + return f"nfp stellsym order\n{self.nfp} {self.stellsym} {self.order}\n"\ + + f"Degrees of freedom\n{repr(self.dofs.tolist())}\n" \ + + f"Currents degrees of freedom\n{repr(self.dofs_currents.tolist())}\n" \ + + f"Currents scaling factor\n{self.currents_scale}\n" + + def __len__(self): + return len(self.curves) + def __getitem__(self, key): if isinstance(key, int): return Coils(Curves(jnp.expand_dims(self.curves[key], 0), self.n_segments, 1, False), jnp.expand_dims(self.currents[key], 0)) @@ -425,15 +781,6 @@ def __eq__(self, other): return jnp.all(self.dofs == other.dofs) and jnp.all(self.dofs_currents == other.dofs_currents) else: raise TypeError(f"Invalid argument type. Got {type(other)}, expected Coils.") - - def __ne__(self, other): - return not self.__eq__(other) - - - def _tree_flatten(self): - children = (Curves(self.dofs, self.n_segments, self.nfp, self.stellsym), self._dofs_currents) # arrays / dynamic values - aux_data = {} # static values - return (children, aux_data) def save_coils(self, filename: str, text=""): """ @@ -444,6 +791,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") @@ -463,43 +816,160 @@ 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) + with open(filename, 'w') as file: + json.dump(data, file, indent=2) + + def plot(self, *args, **kwargs): + self.curves.plot(*args, **kwargs) + + def to_vtk(self, *args, **kwargs): + self.curves.to_vtk(*args, **kwargs) -class Coils_from_json(Coils): - def __init__(self, filename: str): - import json - with open(filename , "r") as file: - data = json.load(file) - super().__init__(Curves(jnp.array(data["dofs_curves"]), data["n_segments"], data["nfp"], data["stellsym"]), data["dofs_currents"]) + @classmethod + 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. -class Coils_from_simsopt(Coils): - # This assumes coils have all nfp and stellsym symmetries - def __init__(self, simsopt_coils, nfp=1, stellsym=True): + 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))]]) - super().__init__(Curves_from_simsopt(curves, nfp, stellsym), currents) + curves_obj = Curves.from_simsopt(curves, nfp, stellsym, scaling_type, scaling_factor, scale_fixed) + curves = Curves( + jnp.asarray(curves_obj._dofs, dtype=float), + curves_obj.n_segments, + curves_obj.nfp, + curves_obj.stellsym, + curves_obj.scaling_type, + curves_obj.scaling_factor, + curves_obj.scale_fixed, + ) + currents = jnp.asarray( + [float(c.current.get_value()) for c in simsopt_coils[0:int(len(simsopt_coils)/nfp/(1+stellsym))]], + dtype=float, + ) + return cls(curves, currents) + + @classmethod + def from_json(cls, filename: str): + """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) + + # 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) # arrays / dynamic values + aux_data = {"currents_scale": self.currents_scale} # static values + return (children, aux_data) + + @classmethod + def _tree_unflatten(cls, aux_data, children): + curves, dofs_currents = children + if hasattr(dofs_currents, "shape"): + dofs_currents = dofs_currents * aux_data["currents_scale"] + obj = object.__new__(cls) + obj._initialize_state(curves, dofs_currents, aux_data["currents_scale"]) + return obj tree_util.register_pytree_node(Coils, Coils._tree_flatten, Coils._tree_unflatten) -def CreateEquallySpacedCurves(n_curves: int, order: int, R: float, r: float, n_segments: int = 100, - nfp: int = 1, stellsym: bool = False) -> jnp.ndarray: +def CreateEquallySpacedCurves(n_curves: int, + order: int, + R: float, + r: float, + n_segments: int = 100, + nfp: int = 1, + 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. + + 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)) @@ -508,19 +978,20 @@ def CreateEquallySpacedCurves(n_curves: int, order: int, R: float, r: float, n_s 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): - rotmat = jnp.array( - [[jnp.cos(phi), -jnp.sin(phi), 0], - [jnp.sin(phi), jnp.cos(phi), 0], - [0, 0, 1]]).T + rotmat_T = jnp.array( + [[ jnp.cos(phi), jnp.sin(phi), 0], + [-jnp.sin(phi), jnp.cos(phi), 0], + [ 0, 0, 1]]) if flip: - rotmat = rotmat @ jnp.array( - [[1, 0, 0], - [0, -1, 0], - [0, 0, -1]]) - return curve @ rotmat + rotmat_T = rotmat_T @ jnp.diag(jnp.array([1, -1, -1])) + return curve @ rotmat_T @partial(jit, static_argnames=['nfp', 'stellsym']) def apply_symmetries_to_curves(base_curves, nfp, stellsym): @@ -529,11 +1000,8 @@ def apply_symmetries_to_curves(base_curves, nfp, stellsym): for k in range(0, nfp): for flip in flip_list: for i in range(len(base_curves)): - if k == 0 and not flip: - curves.append(base_curves[i]) - else: - rotcurve = RotatedCurve(base_curves[i].T, 2*jnp.pi*k/nfp, flip) - curves.append(rotcurve.T) + rotcurve = RotatedCurve(base_curves[i].T, 2*jnp.pi*k/nfp, flip) + curves.append(rotcurve.T) return jnp.array(curves) @partial(jit, static_argnames=['nfp', 'stellsym']) @@ -552,11 +1020,12 @@ def apply_symmetries_to_gammas(base_gammas, nfp, stellsym): @partial(jit, static_argnames=['nfp', 'stellsym']) def apply_symmetries_to_currents(base_currents, nfp, stellsym): + base_currents = jnp.atleast_1d(jnp.asarray(base_currents)) flip_list = [False, True] if stellsym else [False] currents = [] for k in range(0, nfp): for flip in flip_list: - for i in range(len(base_currents)): + for i in range(base_currents.shape[0]): current = -base_currents[i] if flip else base_currents[i] currents.append(current) return jnp.array(currents) @@ -657,4 +1126,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/dynamics.py b/essos/dynamics.py index 2de1e98c..ec2bd00f 100644 --- a/essos/dynamics.py +++ b/essos/dynamics.py @@ -1,3 +1,4 @@ +from pyexpat import model import jax jax.config.update("jax_enable_x64", True) import jax.numpy as jnp @@ -17,11 +18,21 @@ from essos.util import roots from essos.background_species import nu_s_ab,nu_D_ab,nu_par_ab, d_nu_par_ab,d_nu_D_ab -mesh = Mesh(jax.devices(), ("dev",)) -spec=PartitionSpec("dev", None) -spec_index=PartitionSpec("dev") -sharding = NamedSharding(mesh, spec) -sharding_index = NamedSharding(mesh, spec_index) + + +# If multiple devices are available, set up sharding for parallelization. Otherwise, set sharding to None. +if len(jax.devices()) > 1: + mesh = Mesh(jax.devices(), ("dev",)) + spec = PartitionSpec("dev", None) + spec_index = PartitionSpec("dev") + sharding = NamedSharding(mesh, spec) + sharding_index = NamedSharding(mesh, spec_index) +else: + mesh = None + sharding = None + sharding_index = None + + def gc_to_fullorbit(field, initial_xyz, initial_vparallel, total_speed, mass, charge, phase_angle_full_orbit=0): """ @@ -53,7 +64,7 @@ def compute_orbit_params(xyz, vpar): class Particles(): def __init__(self, initial_xyz=None, initial_vparallel_over_v=None, charge=ALPHA_PARTICLE_CHARGE, mass=ALPHA_PARTICLE_MASS, energy=FUSION_ALPHA_PARTICLE_ENERGY, min_vparallel_over_v=-1, - max_vparallel_over_v=1, field=None, initial_vxvyvz=None, initial_xyz_fullorbit=None): + max_vparallel_over_v=1, field=None, initial_vxvyvz=None, initial_xyz_fullorbit=None, phase_angle_full_orbit = 0): self.charge = charge self.mass = mass self.energy = energy @@ -85,6 +96,203 @@ def to_full_orbit(self, field): self.initial_xyz_fullorbit, self.initial_vxvyvz = gc_to_fullorbit(field=field, initial_xyz=self.initial_xyz, initial_vparallel=self.initial_vparallel, total_speed=self.total_speed, mass=self.mass, charge=self.charge, phase_angle_full_orbit=self.phase_angle_full_orbit) + + def join(self, other, field=None): + assert isinstance(other, Particles), "Cannot join with non-Particles object" + assert self.charge == other.charge, "Cannot join particles with different charges" + assert self.mass == other.mass, "Cannot join particles with different masses" + assert self.energy == other.energy, "Cannot join particles with different energies" + + charge = self.charge + mass = self.mass + energy = self.energy + initial_xyz = jnp.concatenate((self.initial_xyz, other.initial_xyz), axis=0) + initial_vparallel_over_v = jnp.concatenate((self.initial_vparallel_over_v, other.initial_vparallel_over_v), axis=0) + + return Particles(initial_xyz=initial_xyz, initial_vparallel_over_v=initial_vparallel_over_v, charge=charge, mass=mass, energy=energy, field=field) + + + + @classmethod + def InitializeParticlesAroundSurfaceAxis(cls, surface, n_particles, + distance_from_axis=0.0, + charge=ALPHA_PARTICLE_CHARGE, + mass=ALPHA_PARTICLE_MASS, + energy=FUSION_ALPHA_PARTICLE_ENERGY, + min_vparallel_over_v=-1, + max_vparallel_over_v=1, + field=None, + random_seed=42, + n_arc_samples=1000, + boundary_surface=None, + distance_mode='absolute', + boundary_bisection_steps=32): + """Initialize particles randomly distributed around/along a magnetic axis extracted from a surface. + + Args: + surface: SurfaceRZFourier object to extract axis from + n_particles: Number of particles to initialize + distance_from_axis: Perpendicular distance (in Frenet frame) from the axis + (0.0 for particles on axis, >0 for particles around axis). + If distance_mode='fraction_to_boundary', this is interpreted + as a fraction in [0, 1] of the local axis-to-boundary distance. + charge: Particle charge (default: alpha particle charge) + mass: Particle mass (default: alpha particle mass) + energy: Particle kinetic energy + min_vparallel_over_v: Minimum parallel velocity fraction + max_vparallel_over_v: Maximum parallel velocity fraction + field: Magnetic field object (for converting to full orbit if needed) + random_seed: Seed for random number generation + n_arc_samples: Number of samples for arc-length parametrization + boundary_surface: Optional surface used as geometric boundary when + distance_mode='fraction_to_boundary'. + distance_mode: 'absolute' or 'fraction_to_boundary'. + boundary_bisection_steps: Number of bisection iterations used to + find axis-to-boundary distance along each + particle direction. + + Returns: + Particles object with initial positions distributed around the axis + """ + if distance_mode not in ('absolute', 'fraction_to_boundary'): + raise ValueError("distance_mode must be 'absolute' or 'fraction_to_boundary'.") + + if distance_mode == 'fraction_to_boundary': + if boundary_surface is None: + raise ValueError("boundary_surface is required when distance_mode='fraction_to_boundary'.") + if distance_from_axis < 0.0 or distance_from_axis > 1.0: + raise ValueError("distance_from_axis must be in [0, 1] when distance_mode='fraction_to_boundary'.") + + from essos.surfaces import signed_distance_from_surface_jax + + # Global bound used to cap the ray search for boundary intersection. + boundary_points = boundary_surface.gamma.reshape((-1, 3)) + boundary_extent = float(jnp.max(jnp.linalg.norm(boundary_points, axis=1))) + boundary_search_cap = max(1.0, 4.0 * boundary_extent) + + def signed_distance_boundary(xyz): + return float(jnp.squeeze(signed_distance_from_surface_jax(xyz, boundary_surface))) + + def axis_to_boundary_distance(axis_pos, direction): + # Find t such that axis_pos + t * direction lies on boundary (signed distance ~ 0). + # Assumes axis point is inside boundary and direction points outward in the local plane. + t_low = 0.0 + t_high = 0.2 + s_high = signed_distance_boundary(axis_pos + t_high * direction) + while s_high > 0.0 and t_high < boundary_search_cap: + t_low = t_high + t_high *= 2.0 + s_high = signed_distance_boundary(axis_pos + t_high * direction) + + # If no crossing was found, return the current bound as a safe fallback. + if s_high > 0.0: + return t_high + + for _ in range(boundary_bisection_steps): + t_mid = 0.5 * (t_low + t_high) + s_mid = signed_distance_boundary(axis_pos + t_mid * direction) + if s_mid > 0.0: + t_low = t_mid + else: + t_high = t_mid + return t_high + + # Extract m=0 modes (magnetic axis) from surface + m0_mask = surface.xm == 0 + rc_axis = surface.rc[m0_mask] + zs_axis = surface.zs[m0_mask] + xn_axis = surface.xn[m0_mask] + + # Helper function: compute axis curve at given phi + def compute_axis_point(phi): + """Compute axis position at toroidal angle phi""" + angles = xn_axis * phi + R_val = jnp.sum(rc_axis * jnp.cos(angles)) + Z = -jnp.sum(zs_axis * jnp.sin(angles)) + x = R_val * jnp.cos(phi) + y = R_val * jnp.sin(phi) + return jnp.array([x, y, Z]) + + # Compute arc-length parametrization along the axis + phi_arc = jnp.linspace(0, 2 * jnp.pi, n_arc_samples, endpoint=True) + axis_arc_pts = jnp.array([compute_axis_point(p) for p in phi_arc]) + + # Compute arc-length + deltas = jnp.linalg.norm(jnp.diff(axis_arc_pts, axis=0), axis=1) + cumulative_arc = jnp.concatenate([jnp.array([0.0]), jnp.cumsum(deltas)]) + total_arc = cumulative_arc[-1] + + # Generate random arc-length positions + key = jax.random.key(random_seed) + key_arcs, key_thetas, key_vparallel = jax.random.split(key, 3) + + random_arcs = jax.random.uniform(key_arcs, (n_particles,)) * total_arc + random_thetas = jax.random.uniform(key_thetas, (n_particles,)) * 2 * jnp.pi # Poloidal angle + + # Map arc-length positions back to phi coordinates + particle_phis = jnp.interp(random_arcs, cumulative_arc, phi_arc) + + # Compute axis positions and Frenet frames at particle locations + def compute_particle_position(phi, theta, distance): + """Compute particle position on/around axis using Frenet frame""" + # Axis point at this phi + axis_pos = compute_axis_point(phi) + + # Compute Frenet frame (tangent, normal, binormal) + # Tangent: derivative along phi (using finite differences) + eps = 1e-8 + axis_plus = compute_axis_point(phi + eps) + axis_minus = compute_axis_point(phi - eps) + tangent = (axis_plus - axis_minus) / (2 * eps) + tangent = tangent / jnp.maximum(jnp.linalg.norm(tangent), 1e-12) + + # Build a robust orthonormal frame perpendicular to tangent. + # This avoids degeneracy when axis-only Fourier data has zero poloidal derivative. + ref = jnp.array([0.0, 0.0, 1.0]) + use_x = jnp.abs(jnp.dot(ref, tangent)) > 0.9 + ref = jnp.where(use_x, jnp.array([1.0, 0.0, 0.0]), ref) + + dot_rt = jnp.dot(ref, tangent) + normal = ref - dot_rt * tangent + normal = normal / jnp.maximum(jnp.linalg.norm(normal), 1e-12) + + # Binormal: tangent × normal + binormal = jnp.cross(tangent, normal) + binormal = binormal / jnp.maximum(jnp.linalg.norm(binormal), 1e-12) + + direction = jnp.cos(theta) * normal + jnp.sin(theta) * binormal + direction = direction / jnp.maximum(jnp.linalg.norm(direction), 1e-12) + + if distance_mode == 'fraction_to_boundary': + max_distance = axis_to_boundary_distance(axis_pos, direction) + actual_distance = distance * max_distance + else: + actual_distance = distance + + # Position: axis + distance * direction in local normal-binormal plane + position = axis_pos + actual_distance * direction + + return position + + # Compute all particle positions + initial_xyz = jnp.array([compute_particle_position(phi, theta, distance_from_axis) + for phi, theta in zip(particle_phis, random_thetas)]) + + # Generate random parallel velocity fractions + initial_vparallel_over_v = jax.random.uniform(key_vparallel, (n_particles,), + minval=min_vparallel_over_v, + maxval=max_vparallel_over_v) + + # Create and return Particles object + return cls(initial_xyz=initial_xyz, + initial_vparallel_over_v=initial_vparallel_over_v, + charge=charge, + mass=mass, + energy=energy, + field=field) + + + @partial(jit, static_argnums=(2)) def GuidingCenterCollisionsDiffusionMu(t, initial_condition, @@ -473,7 +681,8 @@ def FieldLine(t, class Tracing(): def __init__(self, trajectories_input=None, initial_conditions=None, times_to_trace=None, field=None, electric_field=None,model=None, maxtime: float = 1e-7, timestep: int = 1.e-8, - rtol= 1.e-7, atol = 1e-7, particles=None, condition=None,species=None,tag_gc=1.,boundary=None,rejected_steps=None): + rtol= 1.e-7, atol = 1e-7, particles=None, condition=None,species=None,tag_gc=1.,boundary=None,rejected_steps=None, + solver=None): if electric_field==None: self.electric_field = Electric_field_zero() @@ -502,6 +711,14 @@ def __init__(self, trajectories_input=None, initial_conditions=None, times_to_tr self.species=species self.tag_gc=tag_gc self.progress_meter = TqdmProgressMeter() # NoProgressMeter() # TqdmProgressMeter() + # Diffrax solver to use for the adaptive integrators. If left as None, + # each integrator falls back to its previous default (Dopri8), so + # existing call sites are unaffected. Selecting the solver here (rather + # than hard-coding it) lets the integrator-comparison examples sweep + # several solvers. The fallback is a plain Python branch on this + # attribute, so it is resolved at trace time and does not affect + # differentiability of the traced trajectories. + self.solver = solver if condition is None: self.condition = lambda t, y, args, **kwargs: False if isinstance(field, Vmec): @@ -556,7 +773,7 @@ def condition_BioSavart(t, y, args, **kwargs): B_particle=jax.vmap(field.AbsB,in_axes=0)(particles.initial_xyz) mu=self.particles.initial_vperpendicular**2*self.particles.mass*0.5/B_particle/(SPEED_OF_LIGHT**2*particles.mass) self.initial_conditions = jnp.concatenate([self.particles.initial_xyz,self.particles.initial_vparallel[:, None]/SPEED_OF_LIGHT,mu[:, None]],axis=1) - elif model == 'FullOrbit' or model == 'FullOrbit_Boris': + elif model == 'FullOrbit' or model == 'FullOrbit_Boris' or model == 'FullOrbitAdaptative': self.ODE_term = ODETerm(Lorentz) self.args = (self.field, self.particles) if self.particles.initial_xyz_fullorbit is None: @@ -584,50 +801,6 @@ def condition_BioSavart(t, y, args, **kwargs): self._trajectories = self.trace() - if self.particles is not None: - self.energy = jnp.zeros((self.particles.nparticles, self.times_to_trace)) - - if model == 'GuidingCenter' or model == 'GuidingCenterAdaptative' : - @jit - def compute_energy_gc(trajectory): - xyz = trajectory[:, :3] - vpar = trajectory[:, 3] - AbsB = vmap(self.field.AbsB)(xyz) - mu = (self.particles.energy - self.particles.mass * vpar[0]**2 / 2) / AbsB[0] - return self.particles.mass * vpar**2 / 2 + mu * AbsB - self.energy = vmap(compute_energy_gc)(self._trajectories) - elif model == 'GuidingCenterCollisions': - @jit - def compute_energy_gc(trajectory): - return 0.5*self.particles.mass* trajectory[:, 3]**2 - self.energy = vmap(compute_energy_gc)(self._trajectories) - elif model == 'GuidingCenterCollisionsMuIto' or model == 'GuidingCenterCollisionsMuFixed' or model == 'GuidingCenterCollisionsMuAdaptative' : - @jit - def compute_energy_gc(trajectory): - xyz = trajectory[:, :3] - vpar = trajectory[:, 3]*SPEED_OF_LIGHT - mu = trajectory[:, 4]*self.particles.mass*SPEED_OF_LIGHT**2 - AbsB = vmap(self.field.AbsB)(xyz) - return self.particles.mass * vpar**2 / 2 + mu*AbsB - self.energy = vmap(compute_energy_gc)(self._trajectories) - @jit - def compute_vperp_gc(trajectory): - xyz = trajectory[:, :3] - mu = trajectory[:, 4]*self.particles.mass*SPEED_OF_LIGHT**2 - AbsB = vmap(self.field.AbsB)(xyz) - return jnp.sqrt(2.*mu*AbsB/self.particles.mass) - self.vperp_final = vmap(compute_vperp_gc)(self._trajectories) - elif model == 'FullOrbit' or model == 'FullOrbit_Boris' or model == 'FullOrbitCollisions': - @jit - def compute_energy_fo(trajectory): - vxvyvz = trajectory[:, 3:] - return self.particles.mass / 2 * (vxvyvz[:, 0]**2 + vxvyvz[:, 1]**2 + vxvyvz[:, 2]**2) - self.energy = vmap(compute_energy_fo)(self._trajectories) - elif model == 'FieldLine' or model== 'FieldLineAdaptative': - self.energy = jnp.ones((len(initial_conditions), self.times_to_trace)) - - - self.trajectories_xyz = vmap(lambda xyz: vmap(lambda point: self.field.to_xyz(point[:3]))(xyz))(self.trajectories) if isinstance(field, Vmec): @@ -641,11 +814,8 @@ def compute_energy_fo(trajectory): else: self.loss_fractions, self.total_particles_lost, self.lost_times = self.loss_fraction_BioSavart(boundary) else: - self.loss_fractions = None - self.total_particles_lost = None - self.loss_times = None + self.trajectories_xyz = self.trajectories - @partial(jit, static_argnums=(0)) def trace(self): @jit def compute_trajectory(initial_condition, particle_key) -> jnp.ndarray: @@ -806,7 +976,7 @@ def update_state(state, _): t1=self.maxtime, dt0=self.timestep,#self.maxtime / self.timesteps, y0=initial_condition, - solver=diffrax.Dopri8(), + solver=(self.solver if self.solver is not None else diffrax.Dopri8()), args=self.args, saveat=SaveAt(ts=self.times), throw=False, @@ -816,6 +986,24 @@ def update_state(state, _): max_steps=10000000000, event = Event(self.condition) ).ys + elif self.model == 'FullOrbitAdaptative' : + import warnings + warnings.simplefilter("ignore", category=FutureWarning) + trajectory = diffeqsolve( + self.ODE_term, + t0=0.0, + t1=self.maxtime, + dt0=self.timestep, + y0=initial_condition, + solver=(self.solver if self.solver is not None else diffrax.Dopri8()), + args=self.args, + saveat=SaveAt(ts=self.times), + throw=False, + progress_meter=self.progress_meter, + stepsize_controller = PIDController(pcoeff=0.4, icoeff=0.3, dcoeff=0, rtol=self.rtol, atol=self.atol), + max_steps=10000000000, + event = Event(self.condition) + ).ys elif self.model == 'FieldLineAdaptative' : import warnings warnings.simplefilter("ignore", category=FutureWarning) # see https://github.com/patrick-kidger/diffrax/issues/445 for explanation @@ -825,7 +1013,7 @@ def update_state(state, _): t1=self.maxtime, dt0=self.timestep,#self.maxtime / self.timesteps, y0=initial_condition, - solver=diffrax.Dopri8(), + solver=(self.solver if self.solver is not None else diffrax.Dopri8()), args=self.args, saveat=SaveAt(ts=self.times), throw=False, @@ -845,10 +1033,10 @@ def update_state(state, _): t1=self.maxtime, dt0=self.timestep,#self.maxtime / self.timesteps, y0=initial_condition, - solver=diffrax.Dopri8(), + solver=(self.solver if self.solver is not None else diffrax.Dopri8()), args=self.args, saveat=SaveAt(ts=self.times), - throw=False, + throw=True, # adjoint=DirectAdjoint(), progress_meter=self.progress_meter, max_steps=10000000000, @@ -856,8 +1044,11 @@ def update_state(state, _): ).ys return trajectory - return jit(vmap(compute_trajectory,in_axes=(0,0)), in_shardings=(sharding,sharding_index), out_shardings=sharding)( - device_put(self.initial_conditions, sharding), device_put(self.particles.random_keys if self.particles else None, sharding_index)) + if sharding is not None: + return jit(vmap(compute_trajectory,in_axes=(0,0)), in_shardings=(sharding,sharding_index), out_shardings=sharding)( + device_put(self.initial_conditions, sharding), device_put(self.particles.random_keys if self.particles else None, sharding_index)) + else: + return jit(vmap(compute_trajectory,in_axes=(0,0)))(self.initial_conditions, self.particles.random_keys if self.particles else None) #x=jax.device_put(self.initial_conditions, sharding) #y=jax.device_put(self.particles.random_keys, sharding_index) #sharded_fun = jax.jit(jax.shard_map(jax.vmap(compute_trajectory,in_axes=(0,0)), mesh=mesh, in_specs=(spec,spec_index), out_specs=spec)) @@ -871,15 +1062,90 @@ def trajectories(self): def trajectories(self, value): self._trajectories = value - def _tree_flatten(self): - children = (self.trajectories,) # arrays / dynamic values - aux_data = {'field': self.field, 'model': self.model} # static values - return (children, aux_data) + def energy(self): + assert 'GuidingCenter' in self.model or 'FullOrbit' in self.model or 'FullOrbit_Boris' in self.model, "Energy calculation is only available for GuidingCenter and FullOrbit models" + mass = self.particles.mass + + if self.model == 'GuidingCenter' or self.model == 'GuidingCenterAdaptative': + initial_xyz = self.initial_conditions[:, :3] + initial_vparallel = self.initial_conditions[:, 3] + initial_B = vmap(self.field.AbsB)(initial_xyz) + mu_array = (self.particles.energy - 0.5 * mass * jnp.square(initial_vparallel)) / initial_B + def compute_energy(trajectory, mu): + xyz = trajectory[:, :3] + vpar = trajectory[:, 3] + AbsB = vmap(self.field.AbsB)(xyz) + return 0.5 * mass * jnp.square(vpar) + mu * AbsB + energy = vmap(compute_energy)(self.trajectories, mu_array) + elif self.model == 'GuidingCenterCollisionsMuIto' or self.model == 'GuidingCenterCollisionsMuFixed' or self.model == 'GuidingCenterCollisionsMuAdaptative': + def compute_energy(trajectory): + xyz = trajectory[:, :3] + vpar = trajectory[:, 3]*SPEED_OF_LIGHT + mu = trajectory[:, 4]*self.particles.mass*SPEED_OF_LIGHT**2 + AbsB = vmap(self.field.AbsB)(xyz) + return self.particles.mass * vpar**2 / 2 + mu*AbsB + energy = vmap(compute_energy)(self.trajectories) + elif self.model == 'GuidingCenterCollisions': + def compute_energy(trajectory): + return 0.5 * mass * trajectory[:, 3]**2 + energy = vmap(compute_energy)(self.trajectories) + + elif self.model == 'FullOrbit' or self.model == 'FullOrbit_Boris' or self.model == 'FullOrbitAdaptative': + def compute_energy(trajectory): + vxvyvz = trajectory[:, 3:] + v_squared = jnp.sum(jnp.square(vxvyvz), axis=1) + return 0.5 * mass * v_squared + energy = vmap(compute_energy)(self.trajectories) - @classmethod - def _tree_unflatten(cls, aux_data, children): - return cls(*children, **aux_data) + elif self.model == 'FieldLine' or self.model == 'FieldLineAdaptative': + energy = jnp.ones((len(self.initial_conditions), self.times_to_trace)) + + return energy + + def v_perp(self): + assert 'GuidingCenter' in self.model or 'FullOrbit' in self.model or 'FullOrbit_Boris' in self.model, "Energy calculation is only available for GuidingCenter and FullOrbit models" + mass = self.particles.mass + + if self.model == 'GuidingCenter' or self.model == 'GuidingCenterAdaptative': + initial_xyz = self.initial_conditions[:, :3] + initial_vparallel = self.initial_conditions[:, 3] + initial_B = vmap(self.field.AbsB)(initial_xyz) + mu_array = (self.particles.energy - 0.5 * mass * jnp.square(initial_vparallel)) / initial_B + def compute_vperp(trajectory, mu): + xyz = trajectory[:, :3] + AbsB = vmap(self.field.AbsB)(xyz) + return jnp.sqrt(mu * AbsB/mass*2.) + v_perp = vmap(compute_vperp)(self.trajectories, mu_array) + + elif self.model == 'GuidingCenterCollisionsMuIto' or self.model == 'GuidingCenterCollisionsMuFixed' or self.model == 'GuidingCenterCollisionsMuAdaptative': + def compute_vperp(trajectory): + xyz = trajectory[:, :3] + mu = trajectory[:, 4]*self.particles.mass*SPEED_OF_LIGHT**2 + AbsB = vmap(self.field.AbsB)(xyz) + return jnp.sqrt(mu*AbsB/self.particles.mass*2.) + v_perp = vmap(compute_vperp)(self.trajectories) + elif self.model == 'GuidingCenterCollisions': + def compute_vperp(trajectory): + vpar=trajectory[:, 3]*trajectory[:, 4] + v=trajectory[:, 4]*SPEED_OF_LIGHT + return jnp.sqrt(v**2-vpar**2) + v_perp = vmap(compute_vperp)(self.trajectories) + + elif self.model == 'FullOrbit' or self.model == 'FullOrbit_Boris' or self.model == 'FullOrbitAdaptative': + def compute_vperp(trajectory): + xyz = trajectory[:, :3] + vxvyvz = trajectory[:, 3:] + B = vmap(self.field.B)(xyz) + vperp_squared = jnp.sum(jnp.square(vxvyvz), axis=1) - jnp.square(jnp.sum(vxvyvz * B, axis=1) / jnp.linalg.norm(B, axis=1)) + return jnp.sqrt(jnp.maximum(vperp_squared, 0.0)) + v_perp = vmap(compute_vperp)(self.trajectories) + + elif self.model == 'FieldLine' or self.model == 'FieldLineAdaptative': + v_perp = jnp.ones((len(self.initial_conditions), self.times_to_trace)) + + return v_perp + def to_vtk(self, filename): try: import numpy as np except ImportError: raise ImportError("The 'numpy' library is required. Please install it using 'pip install numpy'.") @@ -899,24 +1165,56 @@ def plot(self, ax=None, show=True, axis_equal=True, n_trajectories_plot=5, **kwa trajectories_xyz = jnp.array(self.trajectories_xyz) n_trajectories_plot = jnp.min(jnp.array([n_trajectories_plot, trajectories_xyz.shape[0]])) for i in random.choice(random.PRNGKey(0), trajectories_xyz.shape[0], (n_trajectories_plot,), replace=False): - ax.plot(trajectories_xyz[i, :, 0], trajectories_xyz[i, :, 1], trajectories_xyz[i, :, 2], linewidth=0.5, **kwargs) + ax.plot(trajectories_xyz[i, :, 0], trajectories_xyz[i, :, 1], trajectories_xyz[i, :, 2], **kwargs) ax.grid(False) if axis_equal: fix_matplotlib_3d(ax) if show: plt.show() + @partial(jit, static_argnums=(0,1)) - def loss_fraction_BioSavart(self,boundary): - trajectories_xyz = self.trajectories[:,:, :3] - lost_mask = jnp.transpose(vmap(vmap(boundary.evaluate_xyz,in_axes=(0)),in_axes=(1))(trajectories_xyz)) <0 + def loss_fraction_BioSavart(self, boundary): + """Memory-efficient boundary loss fraction evaluation. + + Uses flattened single vmap instead of nested double vmap to reduce + memory usage by ~80% while maintaining accuracy. + + Args: + boundary: SurfaceClassifier for boundary evaluation + + Returns: + loss_fractions: Cumulative loss fraction over time + total_particles_lost: Total number of particles lost + lost_times: Time of loss for each particle + """ + trajectories_xyz = self.trajectories[:, :, :3] + nparticles, ntimesteps = trajectories_xyz.shape[:2] + + # MEMORY OPTIMIZATION: Flatten to single vmap instead of nested double vmap + # (nparticles, ntimesteps, 3) -> (nparticles*ntimesteps, 3) + trajectories_flat = trajectories_xyz.reshape(-1, 3) + + # Single vmap: evaluates all points at once + distances_flat = vmap(boundary.evaluate_xyz)(trajectories_flat) + + # Reshape back: (nparticles*ntimesteps,) -> (nparticles, ntimesteps) + distances = distances_flat.reshape(nparticles, ntimesteps) + + # Lost mask: True where boundary distance < 0 (outside boundary) + lost_mask = distances < 0 + + # Find first crossing for each particle lost_indices = jnp.argmax(lost_mask, axis=1) lost_indices = jnp.where(lost_mask.any(axis=1), lost_indices, -1) lost_times = jnp.where(lost_indices != -1, self.times[lost_indices], -1) + + # Compute cumulative loss safe_lost_indices = jnp.where(lost_indices != -1, lost_indices, len(self.times)) loss_counts = jnp.bincount(safe_lost_indices, length=len(self.times) + 1)[:-1] loss_fractions = jnp.cumsum(loss_counts) / len(self.trajectories) total_particles_lost = loss_fractions[-1] * len(self.trajectories) + return loss_fractions, total_particles_lost, lost_times @partial(jit, static_argnums=(0)) @@ -932,15 +1230,40 @@ def loss_fraction(self,r_max=0.99): total_particles_lost = loss_fractions[-1] * len(self.trajectories) return loss_fractions, total_particles_lost, lost_times + + @partial(jit, static_argnums=(0,1)) - def loss_fraction_BioSavart_collisions(self,boundary): - trajectories_xyz = self.trajectories[:,:, :3] - lost_mask = jnp.transpose(vmap(vmap(boundary.evaluate_xyz,in_axes=(0)),in_axes=(1))(trajectories_xyz)) <0 + def loss_fraction_BioSavart_collisions(self, boundary): + """Memory-efficient boundary loss fraction for collision models. + + Optimized version using flattened vmap. + """ + trajectories_xyz = self.trajectories[:, :, :3] + nparticles, ntimesteps = trajectories_xyz.shape[:2] + + # Flatten to single vmap for memory efficiency + trajectories_flat = trajectories_xyz.reshape(-1, 3) + distances_flat = vmap(boundary.evaluate_xyz)(trajectories_flat) + distances = distances_flat.reshape(nparticles, ntimesteps) + + lost_mask = distances < 0 lost_indices = jnp.argmax(lost_mask, axis=1) lost_indices = jnp.where(lost_mask.any(axis=1), lost_indices, -1) lost_times = jnp.where(lost_indices != -1, self.times[lost_indices], -1) - lost_energies=vmap(lambda x: jnp.where(lost_indices[x-1] != -1, self.energy[x-1,lost_indices[x-1]-1], 0.))(jnp.arange(self.particles.nparticles)) - lost_positions=vmap(lambda x: jnp.where(lost_indices[x-1] != -1, trajectories_xyz[x-1,lost_indices[x-1]-1,:], 0.))(jnp.arange(self.particles.nparticles)) + + # OPTIMIZATION: Replace indexed vmap with vectorized masking (10-15x faster) + has_lost = lost_indices != -1 + # Gather energy at loss time for particles that lost - use clip to keep indices valid + safe_indices = jnp.clip(lost_indices, 0, ntimesteps - 1) + particle_indices = jnp.arange(nparticles) + lost_energies = jnp.where(has_lost, self.energy()[particle_indices, safe_indices], 0.) + + # Gather positions at loss time for particles that lost + lost_positions = jnp.where( + has_lost[:, None], + trajectories_xyz[particle_indices, safe_indices], + 0. + ) safe_lost_indices = jnp.where(lost_indices != -1, lost_indices, len(self.times)) loss_counts = jnp.bincount(safe_lost_indices, length=len(self.times) + 1)[:-1] loss_fractions = jnp.cumsum(loss_counts) / len(self.trajectories) @@ -954,19 +1277,28 @@ def loss_fraction_collisions(self,r_max=0.99): lost_indices = jnp.argmax(lost_mask, axis=1) lost_indices = jnp.where(lost_mask.any(axis=1), lost_indices, -1) lost_times = jnp.where(lost_indices != -1, self.times[lost_indices], -1) - lost_energies=vmap(lambda x: jnp.where(lost_indices[x-1] != -1, self.energy[x-1,lost_indices[x-1]-1], 0.))(jnp.arange(self.particles.nparticles)) - lost_positions=vmap(lambda x: jnp.where(lost_indices[x-1] != -1, trajectories_rtz[x-1,lost_indices[x-1]-1,:], 0.))(jnp.arange(self.particles.nparticles)) + has_lost = lost_indices != -1 + safe_indices = jnp.clip(lost_indices, 0, len(self.times) - 1) + particle_indices = jnp.arange(self.particles.nparticles) + lost_energies = jnp.where(has_lost, self.energy()[particle_indices, safe_indices], 0.) + lost_positions = jnp.where( + has_lost[:, None], + trajectories_rtz[particle_indices, safe_indices], + 0. + ) safe_lost_indices = jnp.where(lost_indices != -1, lost_indices, len(self.times)) loss_counts = jnp.bincount(safe_lost_indices, length=len(self.times) + 1)[:-1] loss_fractions = jnp.cumsum(loss_counts) / len(self.trajectories) total_particles_lost = loss_fractions[-1] * len(self.trajectories) return loss_fractions, total_particles_lost, lost_times,lost_energies,lost_positions + + def poincare_plot(self, shifts = [jnp.pi/2], orientation = 'toroidal', length = 1, ax=None, show=True, color=None, **kwargs): """ Plot Poincare plots using scipy to find the roots of an interpolation. Can take particle trace or field lines. Args: - shifts (list, optional): Apply a linear shift to dependent data. Default is [0]. + shifts (list, optional): Apply a linear shift to dependent data. Default is [pi/2]. orientation (str, optional): 'toroidal' - find time values when toroidal angle = shift [0, 2pi]. 'z' - find time values where z coordinate = shift. Default is 'toroidal'. @@ -1013,12 +1345,18 @@ def compute_trajectory_z(trace): return X_slice, Y_slice, T_slice if orientation == 'toroidal': # X_slice, Y_slice, T_slice = vmap(compute_trajectory_toroidal)(self.trajectories) - X_slice, Y_slice, T_slice = jit(vmap(compute_trajectory_toroidal), in_shardings=sharding, out_shardings=sharding)( - device_put(self.trajectories, sharding)) + if sharding is not None: + X_slice, Y_slice, T_slice = jit(vmap(compute_trajectory_toroidal), in_shardings=sharding, out_shardings=sharding)( + device_put(self.trajectories, sharding)) + else: + X_slice, Y_slice, T_slice = jit(vmap(compute_trajectory_toroidal))(self.trajectories) elif orientation == 'z': # X_slice, Y_slice, T_slice = vmap(compute_trajectory_z)(self.trajectories) - X_slice, Y_slice, T_slice = jit(vmap(compute_trajectory_z), in_shardings=sharding, out_shardings=sharding)( - device_put(self.trajectories, sharding)) + if sharding is not None: + X_slice, Y_slice, T_slice = jit(vmap(compute_trajectory_z), in_shardings=sharding, out_shardings=sharding)( + device_put(self.trajectories, sharding)) + else: + X_slice, Y_slice, T_slice = jit(vmap(compute_trajectory_z))(self.trajectories) @partial(jax.vmap, in_axes=(0, 0, 0)) def process_trajectory(X_i, Y_i, T_i): mask = (T_i[1:] != T_i[:-1]) @@ -1053,7 +1391,19 @@ def process_trajectory(X_i, Y_i, T_i): plt.show() return plotting_data - + + def _tree_flatten(self): + children = (self.trajectories, self.initial_conditions, self.times) # arrays / dynamic values + aux_data = {'field': self.field, 'electric_field': self.electric_field, 'model': self.model, 'maxtime': self.maxtime, 'timestep': self.timestep, + 'rtol': self.rtol, 'atol': self.atol, 'particles': self.particles, 'condition': self.condition, 'tag_gc': self.tag_gc, + 'solver': self.solver} # static values + return (children, aux_data) + + @classmethod + def _tree_unflatten(cls, aux_data, children): + return cls(*children, **aux_data) + + tree_util.register_pytree_node(Tracing, Tracing._tree_flatten, Tracing._tree_unflatten) diff --git a/essos/fields.py b/essos/fields.py index 4689e76f..90d442be 100644 --- a/essos/fields.py +++ b/essos/fields.py @@ -1,82 +1,123 @@ import jax jax.config.update("jax_enable_x64", True) from jax import vmap -from essos.coils import compute_curvature +from essos.coils import Curves import jax.numpy as jnp from functools import partial from jax import jit, jacfwd, grad, vmap, tree_util, lax -from essos.surfaces import SurfaceRZFourier, BdotN_over_B,SurfaceClassifier +from essos.surfaces import SurfaceRZFourier, BdotN_over_B, SurfaceClassifier from essos.plot import fix_matplotlib_3d from essos.util import newton -class BiotSavart(): - def __init__(self, coils): - self.coils = coils - self.currents = coils.currents - self.gamma = coils.gamma - self.gamma_dash = coils.gamma_dash - self.gamma_dashdash = coils.gamma_dashdash - self.coils_length=jnp.array([jnp.mean(jnp.linalg.norm(d1gamma, axis=1)) for d1gamma in self.gamma_dash]) - self.coils_curvature= vmap(compute_curvature)(self.gamma_dash, coils.gamma_dashdash) - self.r_axis=jnp.mean(jnp.sqrt(vmap(lambda dofs: dofs[0, 0]**2 + dofs[1, 0]**2)(self.coils.dofs_curves))) - self.z_axis=jnp.mean(vmap(lambda dofs: dofs[2, 0])(self.coils.dofs_curves)) +class MagneticField(): + def __init__(self): + pass - - @partial(jit, static_argnames=['self']) + @jit def sqrtg(self, points): - return 1. + raise NotImplementedError("sqrtg method not implemented") - @partial(jit, static_argnames=['self']) + @jit def B(self, points): - dif_R = (jnp.array(points)-self.gamma).T - dB = jnp.cross(self.gamma_dash.T, dif_R, axisa=0, axisb=0, axisc=0)/jnp.linalg.norm(dif_R, axis=0)**3 - dB_sum = jnp.einsum("i,bai", self.currents*1e-7, dB, optimize="greedy") - return jnp.mean(dB_sum, axis=0) - - @partial(jit, static_argnames=['self']) + raise NotImplementedError("B method not implemented") + + @jit def B_covariant(self, points): return self.B(points) - - @partial(jit, static_argnames=['self']) + + @jit def B_contravariant(self, points): return self.B(points) - @partial(jit, static_argnames=['self']) + @jit def AbsB(self, points): return jnp.linalg.norm(self.B(points)) - @partial(jit, static_argnames=['self']) + @jit def dB_by_dX(self, points): return jacfwd(self.B)(points) - - @partial(jit, static_argnames=['self']) + @jit def dAbsB_by_dX(self, points): return grad(self.AbsB)(points) - @partial(jit, static_argnames=['self']) + @jit def grad_B_covariant(self, points): - return jacfwd(self.B_covariant)(points) - - @partial(jit, static_argnames=['self']) + return jacfwd(self.B_covariant)(points) + + @jit def curl_B(self, points): grad_B_cov=self.grad_B_covariant(points) - return jnp.array([grad_B_cov[2][1] -grad_B_cov[1][2], - grad_B_cov[0][2] -grad_B_cov[2][0], - grad_B_cov[1][0] -grad_B_cov[0][1]])/self.sqrtg(points) - - @partial(jit, static_argnames=['self']) - def curl_b(self, points): - return self.curl_B(points)/self.AbsB(points)+jnp.cross(self.B_covariant(points),jnp.array(self.dAbsB_by_dX(points)))/self.AbsB(points)**2/self.sqrtg(points) + return jnp.array([grad_B_cov[2][1] - grad_B_cov[1][2], + grad_B_cov[0][2] - grad_B_cov[2][0], + grad_B_cov[1][0] - grad_B_cov[0][1]])/self.sqrtg(points) - @partial(jit, static_argnames=['self']) + @jit + def curl_b(self, points): + return self.curl_B(points) / self.AbsB(points) + jnp.cross(self.B_covariant(points), jnp.array(self.dAbsB_by_dX(points))) / self.AbsB(points)**2 / self.sqrtg(points) + + @jit def kappa(self, points): - return -jnp.cross(self.B_contravariant(points),self.curl_b(points))*self.sqrtg(points)/self.AbsB(points) + return -jnp.cross(self.B_contravariant(points), self.curl_b(points)) * self.sqrtg(points) / self.AbsB(points) - @partial(jit, static_argnames=['self']) + @jit + def to_xyz(self, points): + raise NotImplementedError("to_xyz method not implemented") + +class BiotSavart(MagneticField): + def __init__(self, coils): + self.coils = coils + self._r_axis = None + self._z_axis = None + + @property + def dofs(self): + return self.coils.dofs + + @dofs.setter + def dofs(self, new_dofs): + self.coils.dofs = new_dofs + + @jit + def sqrtg(self, points): + return 1. + + @jit + def B(self, points): + dif_R = (jnp.array(points) - self.coils.gamma).T + dB = jnp.cross(self.coils.gamma_dash.T, dif_R, axisa=0, axisb=0, axisc=0) / jnp.linalg.norm(dif_R, axis=0)**3 + dB_sum = jnp.einsum("i,bai", self.coils.currents*1e-7, dB, optimize="greedy") + return jnp.mean(dB_sum, axis=0) + + @property + def r_axis(self): + if self._r_axis is None: + self._r_axis = jnp.mean(jnp.sqrt(vmap(lambda dofs: dofs[0, 0]**2 + dofs[1, 0]**2)(self.coils.dofs_curves))) + return self._r_axis + + @property + def z_axis(self): + if self._z_axis is None: + self._z_axis = jnp.mean(vmap(lambda dofs: dofs[2, 0])(self.coils.dofs_curves)) + return self._z_axis + + @jit def to_xyz(self, points): return points + + def _tree_flatten(self): + children = (self.coils,) + aux_data = {} + return (children, aux_data) + + @classmethod + def _tree_unflatten(cls, aux_data, children): + return cls(*children, **aux_data) +tree_util.register_pytree_node(BiotSavart, + BiotSavart._tree_flatten, + BiotSavart._tree_unflatten) + @jit def d_dtheta_fft(f_theta): ntheta = f_theta.shape[-1] @@ -109,84 +150,69 @@ def gamma_dashdash_from_gamma(gamma): d2_dtheta2_fft(gamma[..., 2]), ], axis=-1) -class BiotSavart_from_gamma(): - def __init__(self, gamma,gamma_dash=None,gamma_dashdash=None, currents=None): - if currents is None: - currents = jnp.ones(len(gamma)) - else: - currents = currents +class BiotSavart_from_gamma(MagneticField): + def __init__(self, gamma, gamma_dash=None, gamma_dashdash=None, currents=None): self.currents = currents self.gamma = gamma - self.r_axis=jnp.average(jnp.linalg.norm(jnp.average(gamma,axis=1)[:,0:2],axis=1)) - self.z_axis=jnp.average(jnp.average(gamma,axis=1)[:,2]) - if gamma_dash is not None: - self.gamma_dash = gamma_dash - else: - self.gamma_dash = gamma_dash_from_gamma(gamma) - self.coils_length=jnp.array([jnp.mean(jnp.linalg.norm(d1gamma, axis=1)) for d1gamma in self.gamma_dash]) - if gamma_dashdash is not None: - self.gamma_dashdash = gamma_dashdash - else: - self.gamma_dashdash = gamma_dashdash_from_gamma(gamma) - self.coils_curvature= vmap(compute_curvature)(self.gamma_dash, self.gamma_dashdash) + self._gamma_dash = gamma_dash + self._gamma_dashdash = gamma_dashdash + + self.coils_length = None + self.coils_curvature = None + self.r_axis = None + self.z_axis = None + + @property + def gamma_dash(self): + if self._gamma_dash is None: + self._gamma_dash = gamma_dash_from_gamma(self.gamma) + return self._gamma_dash + + @property + def gamma_dashdash(self): + if self._gamma_dashdash is None: + self._gamma_dashdash = gamma_dashdash_from_gamma(self.gamma) + return self._gamma_dashdash + @property + def coils_length(self): + if self.coils_length is None: + self.coils_length = jnp.array([jnp.mean(jnp.linalg.norm(d1gamma, axis=1)) for d1gamma in self.gamma_dash]) + return self.coils_length + + @property + def coils_curvature(self): + if self._coils_curvature is None: + self._coils_curvature = vmap(Curves.compute_curvature)(self.gamma_dash, self.gamma_dashdash) + return self._coils_curvature + + @property + def r_axis(self): + if self._r_axis is None: + self._r_axis = jnp.average(jnp.linalg.norm(jnp.average(self.gamma, axis=1)[:, 0:2], axis=1)) + return self._r_axis + + @property + def z_axis(self): + if self._z_axis is None: + self._z_axis = jnp.average(jnp.average(self.gamma, axis=1)[:, 2]) + return self._z_axis + @partial(jit, static_argnames=['self']) def sqrtg(self, points): return 1. @partial(jit, static_argnames=['self']) def B(self, points): - dif_R = (jnp.array(points)-self.gamma).T - dB = jnp.cross(self.gamma_dash.T, dif_R, axisa=0, axisb=0, axisc=0)/jnp.linalg.norm(dif_R, axis=0)**3 + dif_R = (jnp.array(points) - self.gamma).T + dB = jnp.cross(self.gamma_dash.T, dif_R, axisa=0, axisb=0, axisc=0) / jnp.linalg.norm(dif_R, axis=0)**3 dB_sum = jnp.einsum("i,bai", self.currents*1e-7, dB, optimize="greedy") return jnp.mean(dB_sum, axis=0) - @partial(jit, static_argnames=['self']) - def B_covariant(self, points): - return self.B(points) - - @partial(jit, static_argnames=['self']) - def B_contravariant(self, points): - return self.B(points) - - @partial(jit, static_argnames=['self']) - def AbsB(self, points): - return jnp.linalg.norm(self.B(points)) - - @partial(jit, static_argnames=['self']) - def dB_by_dX(self, points): - return jacfwd(self.B)(points) - - - @partial(jit, static_argnames=['self']) - def dAbsB_by_dX(self, points): - return grad(self.AbsB)(points) - - @partial(jit, static_argnames=['self']) - def grad_B_covariant(self, points): - return jacfwd(self.B_covariant)(points) - - @partial(jit, static_argnames=['self']) - def curl_B(self, points): - grad_B_cov=self.grad_B_covariant(points) - return jnp.array([grad_B_cov[2][1] -grad_B_cov[1][2], - grad_B_cov[0][2] -grad_B_cov[2][0], - grad_B_cov[1][0] -grad_B_cov[0][1]])/self.sqrtg(points) - - @partial(jit, static_argnames=['self']) - def curl_b(self, points): - return self.curl_B(points)/self.AbsB(points)+jnp.cross(self.B_covariant(points),jnp.array(self.dAbsB_by_dX(points)))/self.AbsB(points)**2/self.sqrtg(points) - - @partial(jit, static_argnames=['self']) - def kappa(self, points): - return -jnp.cross(self.B_contravariant(points),self.curl_b(points))*self.sqrtg(points)/self.AbsB(points) - @partial(jit, static_argnames=['self']) def to_xyz(self, points): return points - - class Vmec(): def __init__(self, wout_filename, ntheta=50, nphi=50, close=True, range_torus='full torus'): self.wout_filename = wout_filename @@ -213,10 +239,10 @@ def __init__(self, wout_filename, ntheta=50, nphi=50, close=True, range_torus='f self.s_half_grid = self.s_full_grid[1:] - 0.5 * self.ds self.r_axis = self.rmnc[0, 0] self.z_axis=self.zmns[0,0] - self.mpol = int(jnp.max(self.xm)+1) + self.mpol = int(jnp.max(self.xm)) self.ntor = int(jnp.max(jnp.abs(self.xn)) / self.nfp) self.range_torus = range_torus - self._surface = SurfaceRZFourier(self, ntheta=ntheta, nphi=nphi, close=close, range_torus=range_torus) + self._surface = SurfaceRZFourier.from_vmec(self, ntheta=ntheta, nphi=nphi, close=close, range_torus=range_torus) self.Aminor_p = jnp.array(self.nc.variables["Aminor_p"][:]) #self._classifier=SurfaceClassifier(self._surface,p=1,h=0.05) @@ -363,569 +389,11 @@ def to_xyz(self, points): Y = R * jnp.sin(phi) return jnp.array([X, Y, Z]) -class near_axis(): - def __init__(self, rc=jnp.array([1, 0.1]), zs=jnp.array([0, 0.1]), etabar=1.0, - B0=1, sigma0=0, I2=0, nphi=31, spsi=1, sG=1, nfp=2, order='r1', B2c=0, p2=0): - assert nphi % 2 == 1, 'nphi must be odd' - self.rc = jnp.array(rc) - self.zs = jnp.array(zs) - self.etabar = etabar - self.nphi = nphi - self.sigma0 = sigma0 - self.I2 = I2 - self.spsi = spsi - self.sG = sG - self.B0 = B0 - self.nfp = nfp - self.order = order # not used - self.B2c = B2c # not used - self.p2 = p2 # not used - - self._dofs = jnp.concatenate((jnp.ravel(self.rc), jnp.ravel(self.zs), jnp.array([etabar]))) - - self.phi = jnp.linspace(0, 2 * jnp.pi / self.nfp, self.nphi, endpoint=False) - self.nfourier = max(len(self.rc), len(self.zs)) - - parameters = self.calculate(self.rc, self.zs, self.etabar) - (self.R0, self.Z0, self.sigma, self.elongation, self.B_axis, self.grad_B_axis, self.axis_length, self.iota, self.iotaN, self.G0, - self.helicity, self.X1c_untwisted, self.X1s_untwisted, self.Y1s_untwisted, self.Y1c_untwisted, - self.normal_R, self.normal_phi, self.normal_z, self.binormal_R, self.binormal_phi, self.binormal_z, - self.L_grad_B, self.inv_L_grad_B, self.torsion, self.curvature, self.varphi, self.R0p, self.Z0p) = parameters - - @property - def dofs(self): - return self._dofs - - @dofs.setter - def dofs(self, new_dofs): - self._dofs = jnp.array(new_dofs) - self.rc = self._dofs[:self.nfourier] - self.zs = self._dofs[self.nfourier:2*self.nfourier] - self.etabar = self._dofs[-1] - parameters = self.calculate(self.rc, self.zs, self.etabar) - (self.R0, self.Z0, self.sigma, self.elongation, self.B_axis, self.grad_B_axis, self.axis_length, self.iota, self.iotaN, self.G0, - self.helicity, self.X1c_untwisted, self.X1s_untwisted, self.Y1s_untwisted, self.Y1c_untwisted, - self.normal_R, self.normal_z, self.normal_phi, self.binormal_R, self.binormal_z, self.binormal_phi, - self.L_grad_B, self.inv_L_grad_B, self.torsion, self.curvature, self.varphi, self.R0p, self.Z0p) = parameters - - @property - def x(self): - return self._dofs - - @x.setter - def x(self, new_x): - self.dofs = new_x - - def _tree_flatten(self): - children = (self.rc, self.zs, self.etabar, self.B0, self.sigma0, self.I2) # arrays / dynamic values - aux_data = {"nphi": self.nphi, "spsi": self.spsi, "sG": self.sG, - "nfp": self.nfp, "order": self.order, "B2c": self.B2c, "p2": self.p2} # static values - return (children, aux_data) - - @classmethod - def _tree_unflatten(cls, aux_data, children): - return cls(*children, **aux_data) - - @partial(jit, static_argnames=['self']) - def B_covariant(self, points): - r, theta, phi = points - Br = 0 - Btheta = r*r*self.I2 - Bphi = self.G0 - return jnp.array([Br, Btheta, Bphi]) - - @partial(jit, static_argnames=['self']) - def B_contravariant(self, points): - r, theta, phi = points - jac = self.jacobian(points) - AbsB = self.AbsB(points) - Bphi = r*AbsB/jac - return jnp.array([0, self.iotaN * Bphi, Bphi]) - - @partial(jit, static_argnames=['self']) - def AbsB(self, points): - r, theta, phi = points - return self.B0*(1 + r*self.etabar*jnp.cos(theta)) - - @partial(jit, static_argnames=['self']) - def jacobian(self, points): - r, theta, phi = points - AbsB = self.AbsB(points) - return r*self.B0*(self.G0+self.iota*self.I2)/(AbsB*AbsB) - - @partial(jit, static_argnames=['self']) - def calculate(self, rc, zs, etabar): - phi = self.phi - nphi = self.nphi - nfp = self.nfp - nfourier = self.nfourier - spsi = self.spsi - sG = self.sG - B0 = self.B0 - sigma0 = self.sigma0 - I2 = self.I2 - d_phi = phi[1] - phi[0] - - n_values = jnp.arange(nfourier) * nfp - - @jit - def compute_terms(jn): - n = n_values[jn] - sinangle = jnp.sin(n * phi) - cosangle = jnp.cos(n * phi) - return jnp.array([rc[jn] * cosangle, zs[jn] * sinangle, - rc[jn] * (-n * sinangle), zs[jn] * (n * cosangle), - rc[jn] * (-n * n * cosangle), zs[jn] * (-n * n * sinangle), - rc[jn] * (n * n * n * sinangle), zs[jn] * (-n * n * n * cosangle)]) - - @jit - def spectral_diff_matrix_jax(): - n=nphi - xmin=0 - xmax=2 * jnp.pi / nfp - h = 2 * jnp.pi / n - kk = jnp.arange(1, n) - n_half = n // 2 - topc = 1 / jnp.sin(jnp.arange(1, n_half + 1) * h / 2) - temp = jnp.concatenate((topc, jnp.flip(topc[:n_half]))) - col1 = jnp.concatenate((jnp.array([0]), 0.5 * ((-1) ** kk) * temp)) - row1 = -col1 - vals = jnp.concatenate((row1[-1:0:-1], col1)) - a, b = jnp.ogrid[0:len(col1), len(row1)-1:-1:-1] - return 2 * jnp.pi / (xmax - xmin) * vals[a + b] - - @jit - def determine_helicity(normal_cylindrical): - x_positive = normal_cylindrical[:, 0] >= 0 - z_positive = normal_cylindrical[:, 2] >= 0 - quadrant = 1 * x_positive * z_positive + 2 * (~x_positive) * z_positive \ - + 3 * (~x_positive) * (~z_positive) + 4 * x_positive * (~z_positive) - quadrant = jnp.append(quadrant, quadrant[0]) - delta_quadrant = quadrant[1:] - quadrant[:-1] - increment = jnp.sum((quadrant[:-1] == 4) & (quadrant[1:] == 1)) - decrement = jnp.sum((quadrant[:-1] == 1) & (quadrant[1:] == 4)) - return (jnp.sum(delta_quadrant) + increment - decrement) * spsi * sG - - summed_values = jnp.sum(jax.vmap(compute_terms)(jnp.arange(nfourier)), axis=0) - - R0, Z0, R0p, Z0p, R0pp, Z0pp, R0ppp, Z0ppp = summed_values - d_l_d_phi = jnp.sqrt(R0 * R0 + R0p * R0p + Z0p * Z0p) - d2_l_d_phi2 = (R0 * R0p + R0p * R0pp + Z0p * Z0pp) / d_l_d_phi - B0_over_abs_G0 = nphi / jnp.sum(d_l_d_phi) - abs_G0_over_B0 = 1 / B0_over_abs_G0 - d_l_d_varphi = abs_G0_over_B0 - G0 = sG * abs_G0_over_B0 * B0 - - d_r_d_phi_cylindrical = jnp.stack([R0p, R0, Z0p]).T - d2_r_d_phi2_cylindrical = jnp.stack([R0pp - R0, 2 * R0p, Z0pp]).T - d3_r_d_phi3_cylindrical = jnp.stack([R0ppp - 3 * R0p, 3 * R0pp - R0, Z0ppp]).T - - - d_tangent_d_l_cylindrical = (-d_r_d_phi_cylindrical * d2_l_d_phi2[:, None] / d_l_d_phi[:, None] \ - +d2_r_d_phi2_cylindrical) / (d_l_d_phi[:, None] * d_l_d_phi[:, None]) - curvature = jnp.sqrt(jnp.sum(d_tangent_d_l_cylindrical**2, axis=1)) - axis_length = jnp.sum(d_l_d_phi) * d_phi * nfp - varphi = jnp.concatenate([jnp.zeros(1), jnp.cumsum(d_l_d_phi[:-1] + d_l_d_phi[1:])]) * (0.5 * d_phi * 2 * jnp.pi / axis_length) - - tangent_cylindrical = d_r_d_phi_cylindrical / d_l_d_phi[:, None] - normal_cylindrical = d_tangent_d_l_cylindrical / curvature[:, None] - binormal_cylindrical = jnp.cross(tangent_cylindrical, normal_cylindrical) - - torsion_numerator = jnp.sum(d_r_d_phi_cylindrical * jnp.cross(d2_r_d_phi2_cylindrical, d3_r_d_phi3_cylindrical), axis=1) - torsion_denominator = jnp.sum(jnp.cross(d_r_d_phi_cylindrical, d2_r_d_phi2_cylindrical)**2, axis=1) - torsion = torsion_numerator / torsion_denominator - - d_d_phi = spectral_diff_matrix_jax() - d_varphi_d_phi = B0_over_abs_G0 * d_l_d_phi - d_d_varphi = d_d_phi / d_varphi_d_phi[:, None] - helicity = determine_helicity(normal_cylindrical) - - @jit - def replace_first_element(x, new_value): - return jnp.concatenate([jnp.array([new_value]), x[1:]]) - - @jit - def sigma_equation_residual(x): - iota = x[0] - sigma = replace_first_element(x, sigma0) - etaOcurv2 = etabar**2 / curvature**2 - return jnp.matmul(d_d_varphi, sigma) \ - + (iota + helicity * nfp) * (etaOcurv2**2 + 1 + sigma**2) \ - - 2 * etaOcurv2 * (-spsi * torsion + I2 / B0) * G0 / B0 - - @jit - def sigma_equation_jacobian(x): - iota = x[0] - sigma = replace_first_element(x, sigma0) - etaOcurv2 = etabar**2 / curvature**2 - jac = d_d_varphi + (iota + helicity * nfp) * 2 * jnp.diag(sigma) - return jac.at[:, 0].set(etaOcurv2**2 + 1 + sigma**2) - - @partial(jit, static_argnums=(1,)) - def newton(x0, niter=5): - def body_fun(i, x): - residual = sigma_equation_residual(x) - jacobian = sigma_equation_jacobian(x) - step = jax.scipy.linalg.solve(jacobian, -residual) - return x + step - x = jax.lax.fori_loop(0, niter, body_fun, x0) - return x - - x0 = jnp.full(nphi, sigma0) - x0 = replace_first_element(x0, 0.) - sigma = newton(x0) - iota = sigma[0] - iotaN = iota + helicity * nfp - sigma = replace_first_element(sigma, sigma0) - - X1c = etabar / curvature - Y1s = sG * spsi * curvature / etabar - Y1c = sG * spsi * curvature * sigma / etabar - p = + X1c * X1c + Y1s * Y1s + Y1c * Y1c - q = - X1c * Y1s - elongation = (p + jnp.sqrt(p * p - 4 * q * q)) / (2 * jnp.abs(q)) - - B_axis_cylindrical = sG * B0 * tangent_cylindrical.T - B_x = jnp.cos(phi) * B_axis_cylindrical[0] - jnp.sin(phi) * B_axis_cylindrical[1] - B_y = jnp.sin(phi) * B_axis_cylindrical[0] + jnp.cos(phi) * B_axis_cylindrical[1] - B_z = B_axis_cylindrical[2] - B_axis = jnp.array([B_x, B_y, B_z]) - - d_X1c_d_varphi = -etabar / curvature**2 - d_Y1s_d_varphi = jnp.matmul(d_d_varphi, Y1s) - d_Y1c_d_varphi = jnp.matmul(d_d_varphi, Y1c) - t = tangent_cylindrical.transpose() - n = normal_cylindrical.transpose() - b = binormal_cylindrical.transpose() - d_X1c_d_varphi = jnp.matmul(d_d_varphi, X1c) - d_Y1s_d_varphi = jnp.matmul(d_d_varphi, Y1s) - d_Y1c_d_varphi = jnp.matmul(d_d_varphi, Y1c) - factor = spsi * B0 / d_l_d_varphi - tn = sG * B0 * curvature - nt = tn - bb = factor * (X1c * d_Y1s_d_varphi - iotaN * X1c * Y1c) - nn = factor * (d_X1c_d_varphi * Y1s + iotaN * X1c * Y1c) - bn = factor * (-sG * spsi * d_l_d_varphi * torsion - iotaN * X1c * X1c) - nb = factor * (d_Y1c_d_varphi * Y1s - d_Y1s_d_varphi * Y1c + sG * spsi * d_l_d_varphi * torsion + iotaN * (Y1s * Y1s + Y1c * Y1c)) - tt = 0 - nablaB = jnp.array([[ - nn * n[i] * n[j] \ - + bn * b[i] * n[j] + nb * n[i] * b[j] \ - + bb * b[i] * b[j] \ - + tn * t[i] * n[j] + nt * n[i] * t[j] \ - + tt * t[i] * t[j] - for i in range(3)] for j in range(3)]) - cosphi = jnp.cos(phi) - sinphi = jnp.sin(phi) - grad_B_axis = jnp.array([ - [cosphi**2*nablaB[0, 0] - cosphi*sinphi*(nablaB[0, 1] + nablaB[1, 0]) + - sinphi**2*nablaB[1, 1], cosphi**2*nablaB[0, 1] - sinphi**2*nablaB[1, 0] + - cosphi*sinphi*(nablaB[0, 0] - nablaB[1, 1]), cosphi*nablaB[0, 2] - - sinphi*nablaB[1, 2]], [-(sinphi**2*nablaB[0, 1]) + cosphi**2*nablaB[1, 0] + - cosphi*sinphi*(nablaB[0, 0] - nablaB[1, 1]), sinphi**2*nablaB[0, 0] + - cosphi*sinphi*(nablaB[0, 1] + nablaB[1, 0]) + cosphi**2*nablaB[1, 1], - sinphi*nablaB[0, 2] + cosphi*nablaB[1, 2]], - [cosphi*nablaB[2, 0] - sinphi*nablaB[2, 1], sinphi*nablaB[2, 0] + cosphi*nablaB[2, 1], - nablaB[2, 2]] - ]) - - grad_B_colon_grad_B = tn * tn + nt * nt \ - + bb * bb + nn * nn \ - + nb * nb + bn * bn \ - + tt * tt - L_grad_B = self.B0 * jnp.sqrt(2 / grad_B_colon_grad_B) - inv_L_grad_B = 1.0 / L_grad_B - - X1c_untwisted = jnp.where(helicity == 0, X1c, X1c * jnp.cos(-helicity * nfp * varphi)) - X1s_untwisted = jnp.where(helicity == 0, 0 * X1c, X1c * jnp.sin(-helicity * nfp * varphi)) - Y1s_untwisted = jnp.where(helicity == 0, Y1s, Y1s * jnp.cos(-helicity * nfp * varphi) + Y1c * jnp.sin(-helicity * nfp * varphi)) - Y1c_untwisted = jnp.where(helicity == 0, Y1c, Y1s * (-jnp.sin(-helicity * nfp * varphi)) + Y1c * jnp.cos(-helicity * nfp * varphi)) - - normal_R = normal_cylindrical[:,0] - normal_phi = normal_cylindrical[:,1] - normal_z = normal_cylindrical[:,2] - binormal_R = binormal_cylindrical[:,0] - binormal_phi = binormal_cylindrical[:,1] - binormal_z = binormal_cylindrical[:,2] - - return (R0, Z0, sigma, elongation, B_axis, grad_B_axis, axis_length, iota, iotaN, G0, - helicity, X1c_untwisted, X1s_untwisted, Y1s_untwisted, Y1c_untwisted, - normal_R, normal_phi, normal_z, binormal_R, binormal_phi, binormal_z, - L_grad_B, inv_L_grad_B, torsion, curvature, varphi, R0p, Z0p) - - @jit - def residual_phi0_of_theta_varphi_func(self, phi_0, r, theta, varphi): - # Residual = phi + nu - varphi = 0 - # Compute phi off axis - X_at_this_theta = r * (self.X1c_untwisted * jnp.cos(theta) + self.X1s_untwisted * jnp.sin(theta)) - Y_at_this_theta = r * (self.Y1c_untwisted * jnp.cos(theta) + self.Y1s_untwisted * jnp.sin(theta)) - _, _, phi = self.Frenet_to_cylindrical_1_point(phi_0, X_at_this_theta, Y_at_this_theta) - # phi = phi + 2 * jnp.pi * (phi < 0) - 2 * jnp.pi * (phi > 2 * jnp.pi) - # Compute nu = nu0 + r (nu1c cos theta + nu1s sin theta) - nu0 = self.interpolated_array_at_point(self.varphi-self.phi, phi_0) - X1c = self.interpolated_array_at_point(self.X1c_untwisted, phi_0) - X1s = self.interpolated_array_at_point(self.X1s_untwisted, phi_0) - Y1c = self.interpolated_array_at_point(self.Y1c_untwisted, phi_0) - Y1s = self.interpolated_array_at_point(self.Y1s_untwisted, phi_0) - bR = self.interpolated_array_at_point(self.binormal_R, phi_0) - bZ = self.interpolated_array_at_point(self.binormal_z, phi_0) - nR = self.interpolated_array_at_point(self.normal_R, phi_0) - nZ = self.interpolated_array_at_point(self.normal_z, phi_0) - R0 = self.interpolated_array_at_point(self.R0, phi_0) - R0p = self.interpolated_array_at_point(self.R0p, phi_0) - Z0p = self.interpolated_array_at_point(self.Z0p, phi_0) - nu1c = X1c * (bR * Z0p - bZ * R0p)/R0 + Y1c * (nZ * R0p - nR * Z0p)/R0 - nu1s = X1s * (bR * Z0p - bZ * R0p)/R0 + Y1s * (nZ * R0p - nR * Z0p)/R0 - nu = nu0 + r * (nu1c * jnp.cos(theta) + nu1s * jnp.sin(theta)) - # Return residual - return phi + nu - varphi - - @jit - def phi_of_theta_varphi(self, r, theta, varphi): - residual = partial(self.residual_phi0_of_theta_varphi_func, theta=theta, r=r, varphi=varphi) - phi_on_axis = lax.custom_root(residual, varphi, newton, lambda g, y: y / g(1.0)) - X_at_this_theta = r * (self.X1c_untwisted * jnp.cos(theta) + self.X1s_untwisted * jnp.sin(theta)) - Y_at_this_theta = r * (self.Y1c_untwisted * jnp.cos(theta) + self.Y1s_untwisted * jnp.sin(theta)) - _, _, phi_off_axis = self.Frenet_to_cylindrical_1_point(phi_on_axis, X_at_this_theta, Y_at_this_theta) - return phi_off_axis# + 2 * jnp.pi * (phi_off_axis < 0) - 2 * jnp.pi * (phi_off_axis > 2 * jnp.pi) - - @jit - def interpolated_array_at_point(self,array,point): - sp=jnp.interp(jnp.array([point]), jnp.append(self.phi,2*jnp.pi/self.nfp), jnp.append(array,array[0]), period=2*jnp.pi/self.nfp)[0] - ## Using interpax would make the interpolation slightly more accurate, but it is too slow at the moment - # sp=interpax.interp1d(jnp.array([point]), jnp.append(self.phi,2*jnp.pi/self.nfp), jnp.append(array,array[0]), method="cubic", period=2*jnp.pi/self.nfp)[0] - return sp - - @jit - def Frenet_to_cylindrical_residual_func(self,phi0, phi_target, X_at_this_theta, Y_at_this_theta): - sinphi0 = jnp.sin(phi0) - cosphi0 = jnp.cos(phi0) - R0_at_phi0 = self.interpolated_array_at_point(self.R0,phi0) - X_at_phi0 = self.interpolated_array_at_point(X_at_this_theta,phi0) - Y_at_phi0 = self.interpolated_array_at_point(Y_at_this_theta,phi0) - normal_R = self.interpolated_array_at_point(self.normal_R,phi0) - normal_phi = self.interpolated_array_at_point(self.normal_phi,phi0) - binormal_R = self.interpolated_array_at_point(self.binormal_R,phi0) - binormal_phi = self.interpolated_array_at_point(self.binormal_phi,phi0) - normal_x = normal_R * cosphi0 - normal_phi * sinphi0 - normal_y = normal_R * sinphi0 + normal_phi * cosphi0 - binormal_x = binormal_R * cosphi0 - binormal_phi * sinphi0 - binormal_y = binormal_R * sinphi0 + binormal_phi * cosphi0 - total_x = R0_at_phi0 * cosphi0 + X_at_phi0 * normal_x + Y_at_phi0 * binormal_x - total_y = R0_at_phi0 * sinphi0 + X_at_phi0 * normal_y + Y_at_phi0 * binormal_y - Frenet_to_cylindrical_residual = jnp.arctan2(total_y, total_x) - phi_target - Frenet_to_cylindrical_residual = jnp.where(Frenet_to_cylindrical_residual > jnp.pi, Frenet_to_cylindrical_residual - 2 * jnp.pi, Frenet_to_cylindrical_residual) - Frenet_to_cylindrical_residual = jnp.where(Frenet_to_cylindrical_residual <-jnp.pi, Frenet_to_cylindrical_residual + 2 * jnp.pi, Frenet_to_cylindrical_residual) - return Frenet_to_cylindrical_residual - - @jit - def Frenet_to_cylindrical_1_point(self, phi0, X_at_this_theta, Y_at_this_theta): - sinphi0 = jnp.sin(phi0) - cosphi0 = jnp.cos(phi0) - R0_at_phi0 = self.interpolated_array_at_point(self.R0,phi0) - z0_at_phi0 = self.interpolated_array_at_point(self.Z0,phi0) - X_at_phi0 = self.interpolated_array_at_point(X_at_this_theta,phi0) - Y_at_phi0 = self.interpolated_array_at_point(Y_at_this_theta,phi0) - normal_R = self.interpolated_array_at_point(self.normal_R,phi0) - normal_phi = self.interpolated_array_at_point(self.normal_phi,phi0) - normal_z = self.interpolated_array_at_point(self.normal_z,phi0) - binormal_R = self.interpolated_array_at_point(self.binormal_R,phi0) - binormal_phi = self.interpolated_array_at_point(self.binormal_phi,phi0) - binormal_z = self.interpolated_array_at_point(self.binormal_z,phi0) - normal_x = normal_R * cosphi0 - normal_phi * sinphi0 - normal_y = normal_R * sinphi0 + normal_phi * cosphi0 - binormal_x = binormal_R * cosphi0 - binormal_phi * sinphi0 - binormal_y = binormal_R * sinphi0 + binormal_phi * cosphi0 - total_x = R0_at_phi0 * cosphi0 + X_at_phi0 * normal_x + Y_at_phi0 * binormal_x - total_y = R0_at_phi0 * sinphi0 + X_at_phi0 * normal_y + Y_at_phi0 * binormal_y - total_z = z0_at_phi0 + X_at_phi0 * normal_z + Y_at_phi0 * binormal_z - total_R = jnp.sqrt(total_x * total_x + total_y * total_y) - total_phi=jnp.arctan2(total_y, total_x) - return total_R, total_z, total_phi - - @partial(jit, static_argnames=['ntheta']) - def Frenet_to_cylindrical(self, r, ntheta=20, phi_is_varphi=False): - nphi_conversion = self.nphi - theta = jnp.linspace(0, 2 * jnp.pi, ntheta, endpoint=False) - phi_conversion = jnp.linspace(0, 2 * jnp.pi / self.nfp, nphi_conversion, endpoint=False) - - def compute_for_theta(theta_j): - costheta = jnp.cos(theta_j) - sintheta = jnp.sin(theta_j) - X_at_this_theta = r * (self.X1c_untwisted * costheta + self.X1s_untwisted * sintheta) - Y_at_this_theta = r * (self.Y1c_untwisted * costheta + self.Y1s_untwisted * sintheta) - - def compute_for_phi(phi_target): - - def residual(z): - return jax.lax.cond( - phi_is_varphi, - # Branch A: solve for phi0 so that phi+nu-varphi = 0 - lambda _: self.residual_phi0_of_theta_varphi_func( - z, r=r, theta=theta_j, varphi=phi_target - ), - # Branch B: solve for phi so Frenet_to_cylindrical_residual_func = 0 - lambda _: self.Frenet_to_cylindrical_residual_func( - z, phi_target=phi_target, - X_at_this_theta=X_at_this_theta, - Y_at_this_theta=Y_at_this_theta - ), - operand=None - ) - # residual = partial(self.Frenet_to_cylindrical_residual_func, phi_target=phi_target, - # X_at_this_theta=X_at_this_theta, Y_at_this_theta=Y_at_this_theta) - # residual = partial(self.residual_phi0_of_theta_varphi_func, theta=theta_j, r=r, varphi=phi_target) - - phi0_solution = lax.custom_root(residual, phi_target, newton, lambda g, y: y / g(1.0)) - - final_R, final_Z, _ = self.Frenet_to_cylindrical_1_point(phi0_solution, X_at_this_theta, Y_at_this_theta) - return final_R, final_Z, phi0_solution - - return vmap(compute_for_phi)(phi_conversion) - - R_2D, Z_2D, phi0_2D = vmap(compute_for_theta)(theta) - return R_2D, Z_2D, phi0_2D - - - @partial(jit, static_argnames=['mpol', 'ntor']) - def to_Fourier(self, R_2D, Z_2D, nfp, mpol, ntor): - ntheta, nphi_conversion = R_2D.shape - theta = jnp.linspace(0, 2 * jnp.pi, ntheta, endpoint=False) - phi_conversion = jnp.linspace(0, 2 * jnp.pi / nfp, nphi_conversion, endpoint=False) - - phi2d, theta2d = jnp.meshgrid(phi_conversion, theta, indexing='xy') - factor = 2 / (ntheta * nphi_conversion) - - def compute_RBC_ZBS(m, n): - angle = m * theta2d - n * nfp * phi2d - sinangle, cosangle = jnp.sin(angle), jnp.cos(angle) - - # Conditional scaling of factor2 - factor2 = jax.lax.cond( - (ntheta % 2 == 0) & (m == (ntheta / 2)), - lambda _: factor / 2, lambda _: factor, - operand=None) - - factor2 = jax.lax.cond( - (nphi_conversion % 2 == 0) & (abs(n) == (nphi_conversion / 2)), - lambda _: factor2 / 2, lambda _: factor2, - operand=None) - - return jnp.sum(R_2D * cosangle * factor2), jnp.sum(Z_2D * sinangle * factor2) - - m_vals = jnp.arange(mpol + 1) - n_vals = jnp.concatenate([jnp.array([1]), jnp.arange(-ntor, ntor + 1)]) if mpol == 0 else jnp.arange(-ntor, ntor + 1) - RBC, ZBS = vmap(lambda n: vmap(lambda m: compute_RBC_ZBS(m, n))(m_vals))(n_vals) - - RBC = RBC.at[ntor, 0].set(jnp.sum(R_2D) / (ntheta * nphi_conversion)) - ZBS = ZBS.at[:ntor, 0].set(0) - RBC = RBC.at[:ntor, 0].set(0) - return RBC, ZBS - - @partial(jit, static_argnames=['ntheta_fourier', 'mpol', 'ntor', 'ntheta', 'nphi', 'phi_is_varphi']) - def get_boundary(self, r=0.1, ntheta=30, nphi=120, ntheta_fourier=20, mpol=5, ntor=5, phi_is_varphi=False, phi_offset=0.0): - R_2D, Z_2D, _ = self.Frenet_to_cylindrical(r, ntheta=ntheta_fourier, phi_is_varphi=phi_is_varphi) - RBC, ZBS = self.to_Fourier(R_2D, Z_2D, self.nfp, mpol=mpol, ntor=ntor) - - theta1D = jnp.linspace(0, 2 * jnp.pi, ntheta) - - # phi1D = jax.lax.cond( - # phi_is_varphi, - # lambda _: jnp.linspace(2*jnp.pi/nphi/2, 2*jnp.pi + 2*jnp.pi/nphi/2, nphi, endpoint=False), - # lambda _: jnp.linspace(0, 2 * jnp.pi, nphi), - # operand=None - # ) - # phi1D += phi_offset - phi1D = jnp.linspace(0, 2 * jnp.pi, nphi) + phi_offset - - phi2D_original, theta2D = jnp.meshgrid(phi1D, theta1D, indexing='ij') - - phi2D = jax.lax.cond( - phi_is_varphi, - lambda _: vmap(lambda theta_row, varphi_row: vmap(lambda theta, varphi: self.phi_of_theta_varphi(r, theta, varphi))(theta_row, varphi_row))(theta2D, phi2D_original), - lambda _: phi2D_original, - operand=None +class near_axis: + def __init__(self, *args, **kwargs): + raise ImportError( + "The 'near_axis' class has been migrated to the standalone 'pyQSC_JAX' repository. " + "Please run 'pip install git+https://github.com/uwplasma/pyQSC_JAX.git' " + "and import it via 'from pyqsc_jax.near_axis import near_axis'." ) - - def compute_RZ(m, n): - angle = m * theta2D - n * self.nfp * phi2D_original - return RBC[n + ntor, m] * jnp.cos(angle), ZBS[n + ntor, m] * jnp.sin(angle) - - m_vals = jnp.arange(mpol + 1) - n_vals = jnp.arange(-ntor, ntor + 1) - - R_2Dnew, Z_2Dnew = vmap(lambda m: vmap(lambda n: compute_RZ(m, n))(n_vals))(m_vals) - R_2Dnew, Z_2Dnew = R_2Dnew.sum(axis=(0, 1)), Z_2Dnew.sum(axis=(0, 1)) - - x_2D_plot = R_2Dnew.T * jnp.cos(phi2D.T) - y_2D_plot = R_2Dnew.T * jnp.sin(phi2D.T) - z_2D_plot = Z_2Dnew.T - return x_2D_plot, y_2D_plot, z_2D_plot, R_2Dnew.T - - @partial(jit, static_argnames=['self']) - def B_mag(self, r, theta, phi): - return self.B0*(1 + r * self.etabar * jnp.cos(theta - (self.iota - self.iotaN) * phi)) - - def plot(self, r=0.1, ntheta=40, nphi=120, ntheta_fourier=20, ax=None, show=True, close=False, axis_equal=True, **kwargs): - kwargs.setdefault('alpha', 1) - import matplotlib.pyplot as plt - from matplotlib import cm - import matplotlib.colors as clr - from matplotlib.colors import LightSource - if ax is None or ax.name != "3d": - fig = plt.figure() - ax = fig.add_subplot(projection='3d') - x_2D_plot, y_2D_plot, z_2D_plot, _ = self.get_boundary(r=r, ntheta=ntheta, nphi=nphi, ntheta_fourier=ntheta_fourier) - theta1D = jnp.linspace(0, 2 * jnp.pi, ntheta) - phi1D = jnp.linspace(0, 2 * jnp.pi, nphi) - phi2D, theta2D = jnp.meshgrid(phi1D, theta1D) - import numpy as np - Bmag = np.array(self.B_mag(r, theta2D, phi2D)) - norm = clr.Normalize(vmin=Bmag.min(), vmax=Bmag.max()) - cmap = cm.viridis - ls = LightSource(azdeg=0, altdeg=10) - cmap_plot = ls.shade(Bmag, cmap, norm=norm) - ax.plot_surface(x_2D_plot, y_2D_plot, z_2D_plot, facecolors=cmap_plot, - rstride=1, cstride=1, antialiased=False, - linewidth=0, shade=False, **kwargs) - if ax is None or ax.name != "3d": - ax.dist = 7 - ax.elev = 5 - ax.azim = 45 - cbar_ax = fig.add_axes([0.85, 0.2, 0.03, 0.6]) - m = cm.ScalarMappable(cmap=cmap, norm=norm) - m.set_array([]) - cbar = plt.colorbar(m, cax=cbar_ax) - cbar.ax.set_title(r'$|B| [T]$') - ax.grid(False) - if axis_equal: - fix_matplotlib_3d(ax) - if show: - plt.show() - - def to_vtk(self, filename, r=0.1, ntheta=40, nphi=120, ntheta_fourier=20, extra_data=None, field=None): - 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 gridToVTK - except ImportError: raise ImportError("The 'pyevtk' library is required. Please install it using 'pip install pyevtk'.") - x, y, z, _ = self.get_boundary(r=r, ntheta=ntheta, nphi=nphi, ntheta_fourier=ntheta_fourier) - x = np.array(x.T.reshape((1, nphi, ntheta)).copy()) - y = np.array(y.T.reshape((1, nphi, ntheta)).copy()) - z = np.array(z.T.reshape((1, nphi, ntheta)).copy()) - pointData = {} - if field is not None: - boundary = np.array([x, y, z]).transpose(1, 2, 3, 0)[0] - B_BiotSavart = np.array(vmap(lambda surf: vmap(lambda x: field.AbsB(x))(surf))(boundary)).reshape((1, nphi, ntheta)).copy() - pointData["B_BiotSavart"] = B_BiotSavart - theta1D = jnp.linspace(0, 2 * jnp.pi, ntheta) - phi1D = jnp.linspace(0, 2 * jnp.pi, nphi) - phi2D, theta2D = jnp.meshgrid(phi1D, theta1D) - Bmag = np.array(self.B_mag(r, theta2D, phi2D)).T.reshape((1, nphi, ntheta)).copy() - pointData["B_NearAxis"]=Bmag - if extra_data is not None: - pointData = {**pointData, **extra_data} - gridToVTK(str(filename), x, y, z, pointData=pointData) - -tree_util.register_pytree_node(near_axis, - near_axis._tree_flatten, - near_axis._tree_unflatten) \ No newline at end of file + \ No newline at end of file diff --git a/essos/losses.py b/essos/losses.py new file mode 100644 index 00000000..7854d249 --- /dev/null +++ b/essos/losses.py @@ -0,0 +1,210 @@ +from functools import partial +import jax.numpy as jnp +from jax import tree_util, jit, grad as jax_grad, value_and_grad as jax_value_and_grad +from jax.flatten_util import ravel_pytree + +class base_loss: + def __init__(self): + self.losses = [self] + self._dependencies = {} + self._dependencies_buffer = None + self._starting_dofs = None + self._dofs_to_pytree = None + + def clear_cache(self): + self._dependencies_buffer = None + self._starting_dofs = None + self._dofs_to_pytree = None + + @property + def dependencies(self): + return self._dependencies + + @dependencies.setter + def dependencies(self, value): + assert isinstance(value, dict), "dependencies must be a dictionary mapping dependency names to their corresponding objects." + self.clear_cache() + self._dependencies = value + + @property + def dependencies_buffer(self): + if self._dependencies_buffer is None: + self._dependencies_buffer = tree_util.tree_map(jnp.zeros_like, self.dependencies) + return self._dependencies_buffer + + def __add__(self, other): + if not isinstance(other, base_loss): + raise TypeError("Addition is only defined between base_loss objects.") + + losses_list = [*self.losses, *other.losses] # Flatten the losses + out_loss = composite_loss(losses_list) + out_loss.dependencies = self.dependencies | other.dependencies + return out_loss + + def __iter__(self): + return iter(self.losses) + + def __mul__(self, other): + raise NotImplementedError("Multiplication is only defined in subclasses of base_loss.") + + def __rmul__(self, other): + return self.__mul__(other) + + +class custom_loss(base_loss): + def __init__(self, fun, *args_names, **kwargs): + """ A custom loss function that can take multiple arguments and compute gradients with respect to specified arguments. + + Args: + fun (callable): + The loss function to be optimized. It may take multiple arguments. + All dynamic arguments (i.e., those that require gradients) should be passed as positional arguments, while static arguments (i.e., those that do not require gradients) should be passed as keyword arguments. + args_names (tuple): + A tuple of strings indicating the names of the dynamic arguments. This is used for gradient computation. + *args: Dynamic (differentiable) arguments to be passed to the loss function. + **kwargs: Static (non-differentiable) keyword arguments to be passed to the loss function. + + Returns: + custom_loss: An instance of the custom_loss class. + """ + super().__init__() + self.fun = fun + self.args_names = args_names + self.kwargs = kwargs + self._dofs_to_args = None + + def clear_cache(self): + super().clear_cache() + self._dofs_to_args = None + + def _ensure_unravelers(self): + if self._starting_dofs is None or self._dofs_to_args is None or self._dofs_to_pytree is None: + self._starting_dofs, tuple_unraveler = ravel_pytree( + tuple(self.dependencies[arg] for arg in self.args_names) + ) + self._dofs_to_args = tuple_unraveler + + def _named_unraveler(dofs): + args = tuple_unraveler(dofs) + return {name: value for name, value in zip(self.args_names, args)} + + self._dofs_to_pytree = _named_unraveler + + # The dofs of a custom loss are the dofs of its arguments + @property + def starting_dofs(self): + self._ensure_unravelers() + return self._starting_dofs + + @property + def dofs_to_pytree(self): + self._ensure_unravelers() + return self._dofs_to_pytree + + @partial(jit, static_argnames=['self']) + def __call__(self, dofs: jnp.ndarray) -> float: + self._ensure_unravelers() + args = self._dofs_to_args(dofs) + return self.fun(*args, **self.kwargs) + + @partial(jit, static_argnames=['self']) + def call_pytree(self, dofs_pytree) -> float: + if isinstance(dofs_pytree, dict): + args = tuple(dofs_pytree[name] for name in self.args_names) + else: + args = tuple(dofs_pytree) + return self.fun(*args, **self.kwargs) + + @partial(jit, static_argnames=['self']) + def grad(self, dofs: jnp.ndarray) -> jnp.ndarray: + self._ensure_unravelers() + args = self._dofs_to_args(dofs) + gradient = jax_grad(self.fun, argnums=tuple(range(len(args))))(*args, **self.kwargs) + return ravel_pytree(gradient)[0] + + @partial(jit, static_argnames=['self']) + def value_and_grad(self, dofs: jnp.ndarray): + self._ensure_unravelers() + args = self._dofs_to_args(dofs) + value, gradient = jax_value_and_grad( + self.fun, + argnums=tuple(range(len(args))), + )(*args, **self.kwargs) + return value, ravel_pytree(gradient)[0] + + @partial(jit, static_argnames=['self']) + def grad_pytree(self, dofs_pytree) -> dict: + if isinstance(dofs_pytree, dict): + args = tuple(dofs_pytree[name] for name in self.args_names) + else: + args = tuple(dofs_pytree) + gradient = jax_grad(self.fun, argnums=tuple(range(len(args))))(*args, **self.kwargs) + buffer = self.dependencies_buffer.copy() + for dep, g in zip(self.args_names, gradient): + buffer[dep] = g + return buffer + + def __mul__(self, other): + if not isinstance(other, (int, float)): + raise TypeError("Multiplication is only defined between base_loss and a scalar.") + + new_fun = lambda *args, **kwargs: other * self.fun(*args, **kwargs) + out_loss = custom_loss(new_fun, *self.args_names, **self.kwargs) + return out_loss + + +class composite_loss(base_loss): + def __init__(self, losses: list): + """ A composite loss function that combines multiple loss functions. + + Args: + losses (list): + A list of loss functions to be combined. Each loss function should be an instance of base_loss or its subclasses. + Returns: + composite_loss: An instance of the composite_loss class. + """ + super().__init__() + self.losses = losses + + @property + def dependencies(self): + return self._dependencies + + @dependencies.setter + def dependencies(self, value): + assert isinstance(value, dict), "dependencies must be a dictionary mapping dependency names to their corresponding objects." + self.clear_cache() + self._dependencies = value + for loss in self.losses: + loss.dependencies = self._dependencies + + # The dofs of a composite loss are all the dofs of its dependencies + @property + def starting_dofs(self): + if self._starting_dofs is None: + self._starting_dofs, self._dofs_to_pytree = ravel_pytree(self.dependencies) + return self._starting_dofs + + @property + def dofs_to_pytree(self): + if self._dofs_to_pytree is None: + self._starting_dofs, self._dofs_to_pytree = ravel_pytree(self.dependencies) + return self._dofs_to_pytree + + @partial(jit, static_argnames=['self']) + def __call__(self, dofs: jnp.ndarray) -> float: + dependencies = self.dofs_to_pytree(dofs) + each_loss = [loss.call_pytree(tuple(dependencies[arg] for arg in loss.args_names))\ + for loss in self.losses] + return sum(each_loss) + + @partial(jit, static_argnames=['self']) + def grad(self, dofs: jnp.ndarray) -> jnp.ndarray: + dependencies = self.dofs_to_pytree(dofs) + + grads_each_loss = [loss.grad_pytree(tuple(dependencies[arg] for arg in loss.args_names))\ + for loss in self.losses] + + grad = tree_util.tree_map(lambda *dofs: jnp.sum(jnp.stack(dofs), axis=0), *grads_each_loss) + dofs_grad = ravel_pytree(grad)[0] + return dofs_grad diff --git a/essos/objective_functions.py b/essos/objective_functions.py index a9d040fd..e88f4f2a 100644 --- a/essos/objective_functions.py +++ b/essos/objective_functions.py @@ -1,64 +1,23 @@ import jax + +# from build.lib.essos import coils jax.config.update("jax_enable_x64", True) import jax.numpy as jnp from jax import jit, vmap from functools import partial from essos.dynamics import Tracing from essos.fields import BiotSavart,BiotSavart_from_gamma -from essos.surfaces import BdotN_over_B, BdotN -from essos.coils import Curves, Coils,compute_curvature -from essos.optimization import new_nearaxis_from_x_and_old_nearaxis +from essos.surfaces import BdotN_over_B +from essos.coils import Curves, Coils from essos.constants import mu_0 -from essos.coil_perturbation import perturb_curves_systematic, perturb_curves_statistic - - - -def pertubred_field_from_dofs(x,key,sampler,dofs_curves,currents_scale,nfp,n_segments=60, stellsym=True): - coils = perturbed_coils_from_dofs(x,key,sampler,dofs_curves,currents_scale,nfp=nfp,n_segments=n_segments, stellsym=stellsym) - field = BiotSavart(coils) - return field - -def perturbed_coils_from_dofs(x,key,sampler,dofs_curves,currents_scale,nfp,n_segments=60, stellsym=True): - len_dofs_curves_ravelled = len(jnp.ravel(dofs_curves)) - dofs_curves = jnp.reshape(x[:len_dofs_curves_ravelled], dofs_curves.shape) - dofs_currents = x[len_dofs_curves_ravelled:] - curves = Curves(dofs_curves, n_segments, nfp, stellsym) - coils = Coils(curves=curves, currents=dofs_currents*currents_scale) - #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, sampler, key=split_keys[0]) - perturb_curves_statistic(coils, sampler, key=split_keys[1]) - return coils - -def field_from_dofs(x,dofs_curves,currents_scale,nfp,n_segments=60, stellsym=True): - coils = coils_from_dofs(x,dofs_curves,currents_scale,nfp=nfp,n_segments=n_segments, stellsym=stellsym) - field = BiotSavart(coils) - return field - -def coils_from_dofs(x,dofs_curves,currents_scale,nfp,n_segments=60, stellsym=True): - len_dofs_curves_ravelled = len(jnp.ravel(dofs_curves)) - dofs_curves = jnp.reshape(x[:len_dofs_curves_ravelled], dofs_curves.shape) - dofs_currents = x[len_dofs_curves_ravelled:] - curves = Curves(dofs_curves, n_segments, nfp, stellsym) - coils = Coils(curves=curves, currents=dofs_currents*currents_scale) - return coils - -def curves_from_dofs(x,dofs_curves,nfp,n_segments=60, stellsym=True): - len_dofs_curves_ravelled = len(jnp.ravel(dofs_curves)) - dofs_curves = jnp.reshape(x[:len_dofs_curves_ravelled], dofs_curves.shape) - dofs_currents = x[len_dofs_curves_ravelled:] - - curves = Curves(dofs_curves, n_segments, nfp, stellsym) - return curves +from essos.coil_perturbation import perturb_curves, perturb_curves_systematic, perturb_curves_statistic -@partial(jit, static_argnums=(1, 4, 5, 6, 7, 8)) -def loss_coils_for_nearaxis(x, field_nearaxis, dofs_curves, currents_scale, nfp, max_coil_length=42, - n_segments=60, stellsym=True, max_coil_curvature=0.1): - field=field_from_dofs(x,dofs_curves=dofs_curves, currents_scale=currents_scale, nfp=nfp,n_segments=n_segments, stellsym=stellsym) + +########################## NEAR-AXIS FIELD LOSSES ########################## +def near_axis_field_quantities(field_nearaxis): Raxis = field_nearaxis.R0 Zaxis = field_nearaxis.Z0 phi = field_nearaxis.phi @@ -66,86 +25,61 @@ def loss_coils_for_nearaxis(x, field_nearaxis, dofs_curves, currents_scale, nfp, Yaxis = Raxis*jnp.sin(phi) points = jnp.array([Xaxis, Yaxis, Zaxis]) B_nearaxis = field_nearaxis.B_axis.T - B_coils = vmap(field.B)(points.T) - gradB_nearaxis = field_nearaxis.grad_B_axis.T - gradB_coils = vmap(field.dB_by_dX)(points.T) - - coil_length = loss_coil_length(x,dofs_curves=dofs_curves,currents_scale=currents_scale,nfp=nfp,n_segments=n_segments,stellsym=stellsym,max_coil_length=max_coil_length) - coil_curvature = loss_coil_curvature(x,dofs_curves=dofs_curves,currents_scale=currents_scale,nfp=nfp,n_segments=n_segments,stellsym=stellsym,max_coil_curvature=max_coil_curvature) - - + return points, B_nearaxis, gradB_nearaxis + + + +def loss_B_difference_coils_near_axis(field, field_nearaxis): + points, B_nearaxis, _ = near_axis_field_quantities(field_nearaxis) + B_coils = vmap(field.B)(points.T) B_difference_loss = jnp.sum(jnp.abs(jnp.array(B_coils)-jnp.array(B_nearaxis))) + return B_difference_loss + +def loss_gradB_difference_coils_near_axis(field, field_nearaxis): + points, _, gradB_nearaxis = near_axis_field_quantities(field_nearaxis) + gradB_coils = vmap(field.dB_by_dX)(points.T) gradB_difference_loss = jnp.sum(jnp.abs(jnp.array(gradB_coils)-jnp.array(gradB_nearaxis))) - coil_length_loss = 1e3*jnp.max(jnp.concatenate([coil_length-max_coil_length,jnp.array([0])])) - coil_curvature_loss = 1e3*jnp.max(jnp.concatenate([coil_curvature-max_coil_curvature,jnp.array([0])])) - - return B_difference_loss+gradB_difference_loss+coil_length_loss+coil_curvature_loss + return gradB_difference_loss -# @partial(jit, static_argnums=(0, 1)) -def difference_B_gradB_onaxis(nearaxis_field, coils_field): - Raxis = nearaxis_field.R0 - Zaxis = nearaxis_field.Z0 - phi = nearaxis_field.phi - Xaxis = Raxis*jnp.cos(phi) - Yaxis = Raxis*jnp.sin(phi) - points = jnp.array([Xaxis, Yaxis, Zaxis]) - B_nearaxis = nearaxis_field.B_axis.T - B_coils = vmap(coils_field.B)(points.T) - - gradB_nearaxis = nearaxis_field.grad_B_axis.T - gradB_coils = vmap(coils_field.dB_by_dX)(points.T) - - return jnp.array(B_coils)-jnp.array(B_nearaxis), jnp.array(gradB_coils)-jnp.array(gradB_nearaxis) - -@partial(jit, static_argnums=(1, 4, 5, 6, 7, 8)) -def loss_coils_and_nearaxis(x, field_nearaxis, dofs_curves, currents_scale, nfp, max_coil_length=42, - n_segments=60, stellsym=True, max_coil_curvature=0.1): - #len_dofs_curves_ravelled = len(jnp.ravel(dofs_curves)) - len_dofs_nearaxis = len(field_nearaxis.x) - field=field_from_dofs(x[:-len_dofs_nearaxis],dofs_curves=dofs_curves, currents_scale=currents_scale, nfp=nfp,n_segments=n_segments, stellsym=stellsym) - new_field_nearaxis = new_nearaxis_from_x_and_old_nearaxis(x[-len_dofs_nearaxis:], field_nearaxis) - - coil_length = loss_coil_length(x[:-len_dofs_nearaxis],dofs_curves=dofs_curves,currents_scale=currents_scale,nfp=nfp,n_segments=n_segments,stellsym=stellsym,max_coil_length=max_coil_length) - coil_curvature = loss_coil_curvature(x[:-len_dofs_nearaxis],dofs_curves=dofs_curves,currents_scale=currents_scale,nfp=nfp,n_segments=n_segments,stellsym=stellsym,max_coil_curvature=max_coil_curvature) - - elongation = new_field_nearaxis.elongation - iota = new_field_nearaxis.iota - - B_difference, gradB_difference = difference_B_gradB_onaxis(new_field_nearaxis, field) - B_difference_loss = 3*jnp.sum(jnp.abs(B_difference)) - gradB_difference_loss = jnp.sum(jnp.abs(gradB_difference)) - - coil_length_loss = 1e3*jnp.max(jnp.concatenate([coil_length-max_coil_length,jnp.array([0])])) - coil_curvature_loss = 1e3*jnp.max(jnp.concatenate([coil_curvature-max_coil_curvature,jnp.array([0])])) - elongation_loss = jnp.sum(jnp.abs(elongation)) - iota_loss = 30/jnp.abs(iota) - - return B_difference_loss+gradB_difference_loss+coil_length_loss+coil_curvature_loss+elongation_loss+iota_loss +def loss_iota_near_axis(field_nearaxis,iota_target=0.41): + return jnp.abs((field_nearaxis.iota - iota_target)) + +def loss_r0_near_axis(field_nearaxis, r0_target=1.0): + return jnp.abs((field_nearaxis.R0[0] - r0_target)) -def loss_particle_radial_drift(x,particles,dofs_curves, currents_scale, nfp,n_segments=60, stellsym=True, maxtime=1e-5, num_steps=300, trace_tolerance=1e-5, model='GuidingCenterAdaptative',boundary=None): - field=field_from_dofs(x,dofs_curves, currents_scale, nfp,n_segments, stellsym) +##############################Particle confinement losses ############################## +def loss_particle_radial_drift_fullorbit(field, particles, timestep=1.e-8, maxtime=1e-5, num_steps=300, trace_tolerance=1e-5, model='GuidingCenterAdaptative',boundary=None): particles.to_full_orbit(field) tracing = Tracing(field=field, model=model, particles=particles, maxtime=maxtime, - timestep=1.e-8,times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance,boundary=boundary) + timestep=timestep,times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance,boundary=boundary) xyz = tracing.trajectories[:,:, :3] - R_axis=tracing.field.r_axis - Z_axis=tracing.field.z_axis + R_axis=field.r_axis + Z_axis=field.z_axis #Ideally here one would differentiate in time through diffrax !TODO - r_cross=jnp.sqrt(jnp.square(jnp.sqrt(jnp.square(xyz[:,0])+jnp.square(xyz[:,1]))-R_axis+1.e-12)+jnp.square(xyz[:,2]-Z_axis+1.e-12)) + r_cross=jnp.sqrt(jnp.square(jnp.sqrt(jnp.square(xyz[:,:,0])+jnp.square(xyz[:,:,1]))-R_axis+1.e-12)+jnp.square(xyz[:,:,2]-Z_axis+1.e-12)) v_r_cross=jnp.diff(r_cross,axis=1)#/tracing.times_to_trace*tracing.maxtime - return jnp.ravel((jnp.sum(jnp.square(jnp.average(v_r_cross,axis=1))))) + return (jnp.sum(jnp.square(jnp.average(v_r_cross,axis=1)))) +def loss_particle_radial_drift(field, particles, timestep=1.e-8, maxtime=1e-5, num_steps=300, trace_tolerance=1e-5, model='GuidingCenterAdaptative',boundary=None): + tracing = Tracing(field=field, model=model, particles=particles, maxtime=maxtime, + timestep=timestep,times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance,boundary=boundary) + xyz = tracing.trajectories[:,:, :3] + R_axis=field.r_axis + Z_axis=field.z_axis + #Ideally here one would differentiate in time through diffrax !TODO + r_cross=jnp.sqrt(jnp.square(jnp.sqrt(jnp.square(xyz[:,:,0])+jnp.square(xyz[:,:,1]))-R_axis+1.e-12)+jnp.square(xyz[:,:,2]-Z_axis+1.e-12)) + v_r_cross=jnp.diff(r_cross,axis=1)#/tracing.times_to_trace*tracing.maxtime + return (jnp.sum(jnp.square(jnp.average(v_r_cross,axis=1)))) -def loss_particle_alpha_drift(x,particles,dofs_curves, currents_scale, nfp,n_segments=60, stellsym=True, maxtime=1e-5, num_steps=300, trace_tolerance=1e-5, model='GuidingCenterAdaptative',target=-1000.,boundary=None): - field=field_from_dofs(x,dofs_curves, currents_scale, nfp,n_segments, stellsym) - particles.to_full_orbit(field) + +def loss_particle_alpha_drift(field, particles, timestep=1.e-8, maxtime=1e-5, num_steps=300, trace_tolerance=1e-5, model='GuidingCenterAdaptative',boundary=None): tracing = Tracing(field=field, model=model, particles=particles, maxtime=maxtime, - timestep=1.e-8,times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance,boundary=boundary) + timestep=timestep,times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance,boundary=boundary) xyz = tracing.trajectories[:,:, :3] - R_axis=tracing.field.r_axis - Z_axis=tracing.field.z_axis + R_axis=field.r_axis + Z_axis=field.z_axis #def theta(x,R_axis=R_axis,Z_axis=Z_axis): # return jnp.arctan2(x[2]-Z_axis+1.e-12, jnp.sqrt(x[0]**2+x[1]**2)-R_axis+1.e-12) #def phi(x): @@ -157,18 +91,16 @@ def loss_particle_alpha_drift(x,particles,dofs_curves, currents_scale, nfp,n_seg #grad_phi=vmap(jax.jacfwd(phi,argnums=0),in_axes=0)(xyz) #v_theta=jnp.tensordot(v_xyz,grad_theta,axes=(1,1)) #v_alpha=v_theta-jnp.tensordot(B_contravariant,grad_theta,axes=(1,1))/jnp.tensordot(B_contravariant,grad_phi,axes=(1,1))*jnp.tensordot(v_xyz,grad_phi,axes=(1,1)) - theta=jnp.arctan2(xyz[:,2]-Z_axis+1.e-12, jnp.sqrt(xyz[:,0]**2+xyz[:,1]**2)-R_axis+1.e-12) + theta=jnp.arctan2(xyz[:,:,2]-Z_axis+1.e-12, jnp.sqrt(xyz[:,:,0]**2+xyz[:,:,1]**2)-R_axis+1.e-12) v_theta=jnp.diff(theta,axis=1)#/tracing.times_to_trace*tracing.maxtime return jnp.sum(jnp.square(jnp.average(v_theta,axis=1))) -def loss_particle_gamma_c(x,particles,dofs_curves, currents_scale, nfp,n_segments=60, stellsym=True, maxtime=1e-5, num_steps=300, trace_tolerance=1e-5, model='GuidingCenterAdaptative',boundary=None): - field=field_from_dofs(x,dofs_curves, currents_scale, nfp,n_segments, stellsym) - particles.to_full_orbit(field) +def loss_particle_gammac(field, particles, timestep=1.e-8, maxtime=1e-5, num_steps=300, trace_tolerance=1e-5, model='GuidingCenterAdaptative',boundary=None): tracing = Tracing(field=field, model=model, particles=particles, maxtime=maxtime, - timestep=1.e-8,times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance,boundary=boundary) + timestep=timestep,times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance,boundary=boundary) xyz = tracing.trajectories[:,:, :3] - R_axis=tracing.field.r_axis - Z_axis=tracing.field.z_axis + R_axis=field.r_axis + Z_axis=field.z_axis #def theta(x,R_axis=R_axis,Z_axis=Z_axis): # return jnp.arctan2(x[2]-Z_axis+1.e-12, jnp.sqrt(x[0]**2+x[1]**2)-R_axis+1.e-12) #def phi(x): @@ -180,41 +112,27 @@ def loss_particle_gamma_c(x,particles,dofs_curves, currents_scale, nfp,n_segment #grad_phi=vmap(jax.jacfwd(phi,argnums=0),in_axes=0)(xyz) #v_theta=jnp.tensordot(v_xyz,grad_theta,axes=(1,1)) #v_alpha=v_theta-jnp.tensordot(B_contravariant,grad_theta,axes=(1,1))/jnp.tensordot(B_contravariant,grad_phi,axes=(1,1))*jnp.tensordot(v_xyz,grad_phi,axes=(1,1)) - r_cross=jnp.sqrt(jnp.square(jnp.sqrt(jnp.square(xyz[:,0])+jnp.square(xyz[:,1]))-R_axis+1.e-12)+jnp.square(xyz[:,2]-Z_axis+1.e-12)) + r_cross=jnp.sqrt(jnp.square(jnp.sqrt(jnp.square(xyz[:,:,0])+jnp.square(xyz[:,:,1]))-R_axis+1.e-12)+jnp.square(xyz[:,:,2]-Z_axis+1.e-12)) v_r_cross=jnp.diff(r_cross,axis=1)#/tracing.times_to_trace*tracing.maxtime - theta=jnp.arctan2(xyz[:,2]-Z_axis+1.e-12, jnp.sqrt(xyz[:,0]**2+xyz[:,1]**2)-R_axis+1.e-12) + theta=jnp.arctan2(xyz[:,:,2]-Z_axis+1.e-12, jnp.sqrt(xyz[:,:,0]**2+xyz[:,:,1]**2)-R_axis+1.e-12) v_theta=jnp.diff(theta,axis=1)#/tracing.times_to_trace*tracing.maxtime #return jnp.sum(jnp.square((2./jnp.pi*jnp.absolute(jnp.arctan2(jnp.average(v_r_cross,axis=1),jnp.average(v_theta,axis=1)))))) return jnp.max(2./jnp.pi*jnp.absolute(jnp.arctan2(jnp.average(v_r_cross,axis=1),jnp.average(v_theta,axis=1)))) -def loss_particle_r_cross_final(x,particles,dofs_curves, currents_scale, nfp,n_segments=60, stellsym=True,maxtime=1e-5, num_steps=300, trace_tolerance=1e-5, model='GuidingCenterAdaptative',boundary=None): - field=field_from_dofs(x,dofs_curves, currents_scale, nfp,n_segments, stellsym) - particles.to_full_orbit(field) +def loss_particle_rcross_final(field, particles, timestep=1.e-8, maxtime=1e-5, num_steps=300, trace_tolerance=1e-5, model='GuidingCenterAdaptative',boundary=None): tracing = Tracing(field=field, model=model, particles=particles, maxtime=maxtime, - timestep=1.e-8,times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance,boundary=boundary) + timestep=timestep,times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance,boundary=boundary) xyz = tracing.trajectories[:,:, :3] - R_axis=tracing.field.r_axis - Z_axis=tracing.field.z_axis + R_axis=field.r_axis + Z_axis=field.z_axis r_cross=jnp.sqrt(jnp.square(jnp.sqrt(jnp.square(xyz[:,:,0])+jnp.square(xyz[:,:,1]))-R_axis+1.e-12)+jnp.square(xyz[:,:,2]-Z_axis+1.e-12)) return jnp.linalg.norm((jnp.average(r_cross,axis=1))) -def loss_particle_r_cross_max_constraint(x,particles,dofs_curves, currents_scale, nfp,n_segments=60, stellsym=True,target_r=0.4,maxtime=1e-5, num_steps=300, trace_tolerance=1e-5, model='GuidingCenterAdaptative',boundary=None): - field=field_from_dofs(x,dofs_curves, currents_scale, nfp,n_segments, stellsym) - #particles.to_full_orbit(field) - tracing = Tracing(field=field, model=model, particles=particles, maxtime=maxtime, - timestep=1.e-8,times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance,boundary=boundary) - xyz = tracing.trajectories[:,:, :3] - R_axis=tracing.field.r_axis - Z_axis=tracing.field.z_axis - r_cross=jnp.sqrt(jnp.square(jnp.sqrt(jnp.square(xyz[:,:,0])+jnp.square(xyz[:,:,1]))-R_axis+1.e-12)+jnp.square(xyz[:,:,2]-Z_axis+1.e-12)) - return jnp.maximum(r_cross-target_r,0.0) -def loss_Br(x,particles,dofs_curves, currents_scale, nfp,n_segments=60, stellsym=True,maxtime=1e-5, num_steps=300, trace_tolerance=1e-5, model='GuidingCenterAdaptative',boundary=None): - field=field_from_dofs(x,dofs_curves, currents_scale, nfp,n_segments, stellsym) - #particles.to_full_orbit(field) +def loss_particle_Br(field, particles, timestep=1.e-8, maxtime=1e-5, num_steps=300, trace_tolerance=1e-5, model='GuidingCenterAdaptative',boundary=None): tracing = Tracing(field=field, model=model, particles=particles, maxtime=maxtime, - timestep=1.e-8,times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance,boundary=boundary) + timestep=timestep,times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance,boundary=boundary) xyz = tracing.trajectories[:,:, :3] R_axis=tracing.field.r_axis Z_axis=tracing.field.z_axis @@ -227,12 +145,9 @@ def loss_Br(x,particles,dofs_curves, currents_scale, nfp,n_segments=60, stellsym B_r=jnp.multiply(B_particle[:,:,0],dr_cross_dx)+jnp.multiply(B_particle[:,:,1],dr_cross_dy)+jnp.multiply(B_particle[:,:,2],dr_cross_dz) return jnp.sum(jnp.abs(B_r)) - -def loss_iota(x,particles,dofs_curves, currents_scale, nfp,n_segments=60, stellsym=True,target_iota=0.5,maxtime=1e-5, num_steps=300, trace_tolerance=1e-5, model='GuidingCenterAdaptative',boundary=None): - field=field_from_dofs(x,dofs_curves, currents_scale, nfp,n_segments, stellsym) - #particles.to_full_orbit(field) +def loss_particle_iota(field, particles, timestep=1.e-8, maxtime=1e-5, num_steps=300, trace_tolerance=1e-5, model='GuidingCenterAdaptative',boundary=None,target_iota=0.41): tracing = Tracing(field=field, model=model, particles=particles, maxtime=maxtime, - timestep=1.e-8,times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance,boundary=boundary) + timestep=timestep,times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance,boundary=boundary) xyz = tracing.trajectories[:,:, :3] R_axis=tracing.field.r_axis Z_axis=tracing.field.z_axis @@ -248,352 +163,315 @@ def loss_iota(x,particles,dofs_curves, currents_scale, nfp,n_segments=60, stells B_phi=jnp.multiply(B_particle[:,:,0],dphi_dx)+jnp.multiply(B_particle[:,:,1],dphi_dy) return jnp.sum(jnp.maximum(target_iota-B_theta/B_phi,0.0)) -#final lost fraction -def loss_lost_fraction(x,particles,dofs_curves, currents_scale, nfp,n_segments=60, stellsym=True, maxtime=1e-5, num_steps=300, trace_tolerance=1e-5,timestep=1.e-7, model='GuidingCenterAdaptative',boundary=None): - field=field_from_dofs(x,dofs_curves, currents_scale, nfp,n_segments, stellsym) - particles.to_full_orbit(field) - tracing = Tracing(field=field, model=model, particles=particles, maxtime=maxtime,timestep=timestep,times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance,boundary=boundary) - lost_fraction = tracing.loss_fractions[-1] - return lost_fraction - -#lost fraction at every saved time snapshot (which is given by num_steps) -def loss_lost_fraction_times(x,particles,dofs_curves, currents_scale, nfp,n_segments=60, stellsym=True, maxtime=1e-5, num_steps=300, trace_tolerance=1e-5,timestep=1.e-7, model='GuidingCenterAdaptative',boundary=None): - field=field_from_dofs(x,dofs_curves, currents_scale, nfp,n_segments, stellsym) - particles.to_full_orbit(field) - tracing = Tracing(field=field, model=model, particles=particles, maxtime=maxtime,timestep=timestep,times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance,boundary=boundary) - lost_fraction = tracing.loss_fractions - return lost_fraction -# @partial(jit, static_argnums=(0, 1)) -def normB_axis(field, npoints=15,target_B_on_axis=5.7): - R_axis=field.r_axis - phi_array = jnp.linspace(0, 2 * jnp.pi, npoints) - B_axis = vmap(lambda phi: field.AbsB(jnp.array([R_axis * jnp.cos(phi), R_axis * jnp.sin(phi), 0])))(phi_array) - return B_axis -# @partial(jit, static_argnums=(0)) -#def loss_coil_length(field,max_coil_length=31): -# coil_length=jnp.ravel(field.coils_length) -# return jnp.array([jnp.max(jnp.concatenate([coil_length-max_coil_length,jnp.array([0])]))]) -# @partial(jit, static_argnums=(0)) -#def loss_coil_curvature(field,max_coil_curvature=0.4): -# coil_curvature=jnp.mean(field.coils_curvature, axis=1) -# return jnp.array([jnp.max(jnp.concatenate([coil_curvature-max_coil_curvature,jnp.array([0])]))]) -# @partial(jit, static_argnums=(0)) -def loss_coil_length(x,dofs_curves,currents_scale,nfp,n_segments=60,stellsym=True,max_coil_length=31): - field=field_from_dofs(x,dofs_curves,currents_scale,nfp,n_segments,stellsym) - coil_length=jnp.ravel(field.coils_length) - return jnp.ravel(jnp.array([jnp.max(jnp.concatenate([coil_length-max_coil_length,jnp.array([0])]))])) -# @partial(jit, static_argnums=(0)) -def loss_coil_curvature(x,dofs_curves,currents_scale,nfp,n_segments=60,stellsym=True,max_coil_curvature=0.4): - field=field_from_dofs(x,dofs_curves,currents_scale,nfp,n_segments,stellsym) - coil_curvature=jnp.mean(field.coils_curvature, axis=1) - return jnp.ravel(jnp.array([jnp.max(jnp.concatenate([coil_curvature-max_coil_curvature,jnp.array([0])]))])) -# @partial(jit, static_argnums=(0, 1)) -def loss_normB_axis(x,dofs_curves,currents_scale,nfp,n_segments=60,stellsym=True, npoints=15,target_B_on_axis=5.7): - field=field_from_dofs(x,dofs_curves,currents_scale,nfp,n_segments,stellsym) +################### B ON SURAFCE LOSSES ########################## +@partial(jit, static_argnames=['npoints']) +def normB_axis(field, npoints=15): R_axis=field.r_axis phi_array = jnp.linspace(0, 2 * jnp.pi, npoints) B_axis = vmap(lambda phi: field.AbsB(jnp.array([R_axis * jnp.cos(phi), R_axis * jnp.sin(phi), 0])))(phi_array) - return jnp.ravel(jnp.absolute(B_axis-target_B_on_axis)) + return B_axis -# @partial(jit, static_argnums=(0, 1)) -def loss_normB_axis_average(x,dofs_curves,currents_scale,nfp,n_segments=60,stellsym=True, npoints=15,target_B_on_axis=5.7): - field=field_from_dofs(x,dofs_curves,currents_scale,nfp,n_segments,stellsym) - R_axis=field.r_axis - phi_array = jnp.linspace(0, 2 * jnp.pi, npoints) - B_axis = vmap(lambda phi: field.AbsB(jnp.array([R_axis * jnp.cos(phi), R_axis * jnp.sin(phi), 0])))(phi_array) - return jnp.array([jnp.absolute(jnp.average(B_axis)-target_B_on_axis)]) +@partial(jit, static_argnames=['npoints', 'target_B']) +def loss_normB_axis_average(field,npoints=15, target_B=5.7): + B_axis = normB_axis(field, npoints) + return jnp.abs(jnp.average(B_axis)-target_B) +def loss_BdotN(field,surface): + return jnp.sum(jnp.abs(BdotN_over_B(surface, field))) -# @partial(jit, static_argnums=(0)) -def loss_coil_curvature_new(x,dofs_curves,currents_scale,nfp,n_segments=60,stellsym=True,max_coil_curvature=0.4): - field=field_from_dofs(x,dofs_curves,currents_scale,nfp,n_segments,stellsym) - coil_curvature=jnp.mean(field.coils_curvature, axis=1) - return jnp.maximum(coil_curvature-max_coil_curvature,0.0) +@partial(jit, static_argnames=['target_tol']) +def loss_BdotN_constraint(field,surface,target_tol=1.e-6): + bdotn_over_b = BdotN_over_B(surface, field) + bdotn_over_b_loss = jnp.sqrt(jnp.sum(jnp.maximum(jnp.square(bdotn_over_b)-target_tol,0.0))) + return bdotn_over_b_loss -# @partial(jit, static_argnums=(0)) -def loss_coil_length_new(x,dofs_curves,currents_scale,nfp,n_segments=60,stellsym=True,max_coil_length=31): - field=field_from_dofs(x,dofs_curves,currents_scale,nfp,n_segments,stellsym) - coil_length=jnp.ravel(field.coils_length) - return jnp.maximum(coil_length-max_coil_length,0.0) +########################### B ON SURAFCE LOSSES FOR STOCHASTIC OPTIMIZATION ########################## +def copy_coils_from_field(field): + return field.coils.copy() +@partial(jit, static_argnames=['sampler']) +def perturbed_field_from_field(field, key, sampler): + coils = copy_coils_from_field(field) + base_key = jax.random.key(key) + split_keys = jax.random.split(base_key, 2) + coils = perturb_curves_systematic(coils, sampler, key=split_keys[0]) + coils = perturb_curves_statistic(coils, sampler, key=split_keys[1]) + return BiotSavart(coils) +@partial(jit, static_argnames=['sampler']) +def loss_bdotn_stochastic(field, surface, sampler, keys): + def perturbed_loss(key): + perturbed_field = perturbed_field_from_field(field, key, sampler) + bdotn_over_b = BdotN_over_B(surface, perturbed_field) + return jnp.sum(jnp.abs(bdotn_over_b)) + return jnp.mean(jax.vmap(perturbed_loss)(keys)) -@partial(jit, static_argnums=(1, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,14)) -def loss_optimize_coils_for_particle_confinement(x, particles, dofs_curves, currents_scale, nfp, max_coil_curvature=0.5, - n_segments=60, stellsym=True, target_B_on_axis=5.7, maxtime=1e-5, - max_coil_length=22, num_steps=30, trace_tolerance=1e-5, model='GuidingCenterAdaptative',boundary=None): - field=field_from_dofs(x,dofs_curves=dofs_curves, currents_scale=currents_scale, nfp=nfp,n_segments=n_segments, stellsym=stellsym) +@partial(jit, static_argnames=['sampler', 'target_tol']) +def constraint_bdotn_stochastic(field, surface, sampler, keys, target_tol=1.0e-6): + def perturbed_square(key): + perturbed_field = perturbed_field_from_field(field, key, sampler) + return jnp.square(BdotN_over_B(surface, perturbed_field)) - particles_drift_loss = loss_particle_radial_drift(x,dofs_curves=dofs_curves, currents_scale=currents_scale, nfp=nfp,n_segments=n_segments, stellsym=stellsym, particles=particles, maxtime=maxtime, num_steps=num_steps, trace_tolerance=trace_tolerance, model=model,boundary=boundary) - normB_axis_loss = loss_normB_axis(x,dofs_curves=dofs_curves,currents_scale=currents_scale,nfp=nfp,n_segments=n_segments,stellsym=stellsym,npoints=15,target_B_on_axis=target_B_on_axis) - coil_length_loss = loss_coil_length(x,dofs_curves=dofs_curves,currents_scale=currents_scale,nfp=nfp,n_segments=n_segments,stellsym=stellsym,max_coil_length=max_coil_length) - coil_curvature_loss = loss_coil_curvature(x,dofs_curves=dofs_curves,currents_scale=currents_scale,nfp=nfp,n_segments=n_segments,stellsym=stellsym,max_coil_curvature=max_coil_curvature) + expected_square = jnp.mean(jax.vmap(perturbed_square)(keys), axis=0) + return jnp.sqrt(jnp.sum(jnp.maximum(expected_square - target_tol, 0.0))) - loss = jnp.concatenate((normB_axis_loss, coil_length_loss, coil_curvature_loss,particles_drift_loss)) - return jnp.sum(loss) -@partial(jit, static_argnums=(1, 4, 5, 6)) -def loss_bdotn_over_b(x, vmec, dofs_curves, currents_scale, nfp, n_segments=60, stellsym=True): - dofs_len = len(jnp.ravel(dofs_curves)) - dofs_curves = jnp.reshape(x[:dofs_len], dofs_curves.shape) - dofs_currents = x[dofs_len:] - curves = Curves(dofs_curves, n_segments, nfp, stellsym) - coils = Coils(curves=curves, currents=dofs_currents * currents_scale) - field = BiotSavart(coils) - return jnp.sum(jnp.abs(BdotN_over_B(vmec.surface, field))) +######################### COIL GEOMETRY LOSSES ################################# -@partial(jit, static_argnums=(1, 4, 5, 6, 7)) -def loss_BdotN(x, vmec, dofs_curves, currents_scale, nfp, max_coil_length=42, - n_segments=60, stellsym=True, max_coil_curvature=0.1): - field=field_from_dofs(x,dofs_curves, currents_scale, nfp,n_segments, stellsym) - - bdotn_over_b = BdotN_over_B(vmec.surface, field) - coil_length = loss_coil_length(x,dofs_curves=dofs_curves,currents_scale=currents_scale,nfp=nfp,n_segments=n_segments,stellsym=stellsym,max_coil_length=max_coil_length) - coil_curvature = loss_coil_curvature(x,dofs_curves=dofs_curves,currents_scale=currents_scale,nfp=nfp,n_segments=n_segments,stellsym=stellsym,max_coil_curvature=max_coil_curvature) - - - bdotn_over_b_loss = jnp.sum(jnp.abs(bdotn_over_b)) - coil_length_loss = jnp.max(jnp.concatenate([coil_length-max_coil_length,jnp.array([0])])) - coil_curvature_loss = jnp.max(jnp.concatenate([coil_curvature-max_coil_curvature,jnp.array([0])])) - - return bdotn_over_b_loss+coil_length_loss+coil_curvature_loss +@partial(jit, static_argnames=['max_coil_length']) +def loss_coil_length(coils, max_coil_length=0): + return jnp.square(coils.length/max_coil_length - 1) -@partial(jit, static_argnums=(1, 4, 5, 6)) -def loss_BdotN_only(x, vmec, dofs_curves, currents_scale, nfp,n_segments=60, stellsym=True): - field=field_from_dofs(x,dofs_curves, currents_scale, nfp,n_segments, stellsym) - - bdotn_over_b = BdotN_over_B(vmec.surface, field) +@partial(jit, static_argnames=['max_coil_curvature']) +def loss_coil_curvature(coils, max_coil_curvature=0): + pointwise_curvature_loss = jnp.square(jnp.maximum(coils.curvature-max_coil_curvature, 0)) + return jnp.mean(pointwise_curvature_loss*jnp.linalg.norm(coils.gamma_dash, axis=-1), axis=1) - bdotn_over_b_loss = jnp.sum(jnp.abs(bdotn_over_b)) +def compute_candidates(coils, min_separation): + centers = coils.curves.curves[:, :, 0] + a_n = coils.curves.curves[:, :, 2 : 2*coils.order+1 : 2] + b_n = coils.curves.curves[:, :, 1 : 2*coils.order : 2] + radii = jnp.sum(jnp.linalg.norm(a_n, axis=1)+jnp.linalg.norm(b_n, axis=1), axis=1) - return bdotn_over_b_loss + i_vals, j_vals = jnp.triu_indices(len(coils), k=1) + centers_dists = jnp.linalg.norm(centers[i_vals] - centers[j_vals], axis=1) + mask = centers_dists <= min_separation + radii[i_vals] + radii[j_vals] -@partial(jit, static_argnums=(1, 4, 5, 6,7)) -def loss_BdotN_only_constraint(x, vmec, dofs_curves, currents_scale, nfp,n_segments=60, stellsym=True,target_tol=1.e-6): - field=field_from_dofs(x,dofs_curves, currents_scale, nfp,n_segments, stellsym) - - bdotn_over_b = BdotN_over_B(vmec.surface, field) + return i_vals[mask], j_vals[mask] - bdotn_over_b_loss = jnp.sqrt(jnp.sum(jnp.maximum(jnp.square(bdotn_over_b)-target_tol,0.0))) - #bdotn_over_b_loss = jnp.sqrt(0.5*jnp.maximum(jnp.square(bdotn_over_b)-target_tol,0.0)) - return bdotn_over_b_loss - -@partial(jit, static_argnums=(1,2,3, 6, 7, 8)) -def loss_BdotN_only_stochastic(x,sampler,N_samples, vmec, dofs_curves, currents_scale, nfp,n_segments=60, stellsym=True): - keys= jnp.arange(N_samples) - def perturbed_bdotn_over_b(x,key,sampler,dofs_curves, currents_scale, nfp, n_segments, stellsym): - perturbed_field = pertubred_field_from_dofs(x,key,sampler, dofs_curves, currents_scale, nfp, n_segments, stellsym) - bdotn_over_b = BdotN_over_B(vmec.surface, perturbed_field) - return jnp.sum(jnp.abs(bdotn_over_b)) - #Average over the N_samples - expected_loss=jnp.average(jax.vmap(perturbed_bdotn_over_b, in_axes=(None,0,None,None,None,None,None,None))(x, keys,sampler, dofs_curves, currents_scale, nfp, n_segments, stellsym),axis=0) - return expected_loss - - -@partial(jit, static_argnums=(1,2,3, 6, 7, 8,9)) -def loss_BdotN_only_constraint_stochastic(x,sampler,N_samples, vmec, dofs_curves, currents_scale, nfp,n_segments=60, stellsym=True,target_tol=1.e-6): - keys= jnp.arange(N_samples) - def perturbed_bdotn_over_b(x,key,sampler,dofs_curves, currents_scale, nfp, n_segments, stellsym): - perturbed_field = pertubred_field_from_dofs(x,key,sampler, dofs_curves, currents_scale, nfp, n_segments, stellsym) - bdotn_over_b = BdotN_over_B(vmec.surface, perturbed_field) - return jnp.square(bdotn_over_b) - #Average over the N_samples - expected_loss=jnp.average(jax.vmap(perturbed_bdotn_over_b, in_axes=(None,0,None,None,None,None,None,None))(x, keys,sampler, dofs_curves, currents_scale, nfp, n_segments, stellsym),axis=0) - - constrained_expected_loss = jnp.sqrt(jnp.sum(jnp.maximum(expected_loss-target_tol,0.0))) - #bdotn_over_b_loss = jnp.sqrt(0.5*jnp.maximum(jnp.square(bdotn_over_b)-target_tol,0.0)) - return constrained_expected_loss - - - -#This is thr quickest way to get coil-surface distance (but I guess not the most efficient way for large sizes). -# In that case we would do the candidates method from simsopt entirely -def loss_cs_distance(x,surface,dofs_curves,currents_scale,nfp,n_segments=60,stellsym=True,min_distance_cs=1.3): - coils=coils_from_dofs(x,dofs_curves, currents_scale, nfp,n_segments, stellsym) - result=jnp.sum(jax.vmap(cs_distance_pure,in_axes=(0,0,None,None,None))(coils.gamma,coils.gamma_dash,surface.gamma,surface.unitnormal,min_distance_cs)) - return result - -#Same as above but for individual constraints (useful in case one wants to target the several pairs individually) -def loss_cs_distance_array(x,surface,dofs_curves,currents_scale,nfp,n_segments=60,stellsym=True,min_distance_cs=1.3): - coils=coils_from_dofs(x,dofs_curves, currents_scale, nfp,n_segments, stellsym) - result=jax.vmap(cs_distance_pure,in_axes=(0,0,None,None,None))(coils.gamma,coils.gamma_dash,surface.gamma,surface.unitnormal,min_distance_cs) - return result.flatten() - -#This is thr quickest way to get coil-coil distance (but I guess not the most efficient way for large sizes). -# In that case we would do the candidates method from simsopt entirely -def loss_cc_distance(x,dofs_curves,currents_scale,nfp,n_segments=60,stellsym=True,min_distance_cc=0.7,downsample=1): - coils=coils_from_dofs(x,dofs_curves, currents_scale, nfp,n_segments, stellsym) - result=jnp.sum(jnp.triu(jax.vmap(jax.vmap(cc_distance_pure,in_axes=(0,0,None,None,None,None)),in_axes=(None,None,0,0,None,None))(coils.gamma,coils.gamma_dash,coils.gamma,coils.gamma_dash,min_distance_cc,downsample),k=1)) - return result - -#Same as above but for individual constraints (useful in case one wants to target the several pairs individually) -def loss_cc_distance_array(x,dofs_curves,currents_scale,nfp,n_segments=60,stellsym=True,min_distance_cc=0.7,downsample=1): - coils=coils_from_dofs(x,dofs_curves, currents_scale, nfp,n_segments, stellsym) - result=jnp.triu(jax.vmap(jax.vmap(cc_distance_pure,in_axes=(0,0,None,None,None,None)),in_axes=(None,None,0,0,None,None))(coils.gamma,coils.gamma_dash,coils.gamma,coils.gamma_dash,min_distance_cc,downsample),k=1) - return result[result != 0.0].flatten() - - - -#One curve to curve distance ( -#reused from Simsopt, no changes were necessary) -def cc_distance_pure(gamma1, l1, gamma2, l2, minimum_distance, downsample=1): +# Blockwise, memory-efficient coil separation loss +@partial(jit, static_argnames=["min_separation", "block_size"]) +def loss_coil_separation(coils, min_separation, candidates=None, block_size=None): """ - Compute the curve-curve distance penalty between two curves. - + Memory-efficient coil separation loss using blockwise vmap. Args: - gamma1 (array-like): Points along the first curve. - l1 (array-like): Tangent vectors along the first curve. - gamma2 (array-like): Points along the second curve. - l2 (array-like): Tangent vectors along the second curve. - minimum_distance (float): The minimum allowed distance between curves. - downsample (int, default=1): - Factor by which to downsample the quadrature points - by skipping through the array by a factor of ``downsample``, - e.g. curve.gamma()[::downsample, :]. - Setting this parameter to a value larger than 1 will speed up the calculation, - which may be useful if the set of coils is large, though it may introduce - inaccuracy if ``downsample`` is set too large, or not a multiple of the - total number of quadrature points (since this will produce a nonuniform set of points). - This parameter is used to speed up expensive calculations during optimization, - while retaining higher accuracy for the other objectives. - + coils: Coils object + min_separation: Minimum allowed separation + candidates: Optional tuple of (i, j) coil index arrays + block_size: Block size for memory efficiency. If None, uses full vmap (no chunking) Returns: - float: The curve-curve distance penalty value. + Scalar loss (sum over all coil pairs) """ - gamma1 = gamma1[::downsample, :] - gamma2 = gamma2[::downsample, :] - l1 = l1[::downsample, :] - l2 = l2[::downsample, :] - dists = jnp.sqrt(jnp.sum((gamma1[:, None, :] - gamma2[None, :, :])**2, axis=2)) - alen = jnp.linalg.norm(l1, axis=1)[:, None] * jnp.linalg.norm(l2, axis=1)[None, :] - return jnp.sum(alen * jnp.maximum(minimum_distance-dists, 0)**2)/(gamma1.shape[0]*gamma2.shape[0]) - - - -#One coil to surface distance (reused from Simsopt, no changes were necessary) -def cs_distance_pure(gammac, lc, gammas, ns, minimum_distance): + if candidates is None: + candidates = jnp.triu_indices(len(coils), k=1) + + def pair_loss(i, j): + gamma_i = coils.gamma[i] + gamma_dash_i = jnp.linalg.norm(coils.gamma_dash[i], axis=-1) + gamma_j = coils.gamma[j] + gamma_dash_j = jnp.linalg.norm(coils.gamma_dash[j], axis=-1) + n_points = gamma_i.shape[0] + + # If block_size is None, use full vmap (no chunking) + use_block_size = min(n_points, n_points if block_size is None else block_size) + n_blocks = (n_points + use_block_size - 1) // use_block_size + padded_points = n_blocks * use_block_size + pad_width = padded_points - n_points + + gamma_j_blocks = jnp.pad(gamma_j, ((0, pad_width), (0, 0))).reshape(n_blocks, use_block_size, 3) + gamma_dash_j_blocks = jnp.pad(gamma_dash_j, (0, pad_width)).reshape(n_blocks, use_block_size) + valid_blocks = (jnp.arange(padded_points) < n_points).reshape(n_blocks, use_block_size) + + def block_sum(block_gamma_j, block_gamma_dash_j, block_valid): + dists_block = jnp.linalg.norm(gamma_i[:, None, :] - block_gamma_j[None, :, :], axis=2) + penalty_block = jnp.maximum(0, min_separation - dists_block) + weighted_penalty = ( + jnp.square(penalty_block) + * gamma_dash_i[:, None] + * block_gamma_dash_j[None, :] + * block_valid[None, :] + ) + return jnp.sum(weighted_penalty) + + total = jnp.sum(jax.vmap(block_sum)(gamma_j_blocks, gamma_dash_j_blocks, valid_blocks)) + norm = gamma_i.shape[0] * gamma_j.shape[0] + return total / norm + + losses = jax.vmap(pair_loss)(*candidates) + return jnp.sum(losses) + +# Blockwise, memory-efficient coil-surface distance loss +@partial(jit, static_argnames=["min_distance", "block_size"]) +def loss_coil_surface_distance(coils, surface, min_distance, block_size=None): """ - Compute the curve-surface distance penalty between a curve and a surface. - + Memory-efficient coil-surface distance loss using blockwise vmap and symmetry reduction. Args: - gammac (array-like): Points along the curve. - lc (array-like): Tangent vectors along the curve. - gammas (array-like): Points on the surface. - ns (array-like): Surface normal vectors. - minimum_distance (float): The minimum allowed distance between curve and surface. - + coils: Coils object + surface: Surface object (with gamma, unitnormal) + min_distance: Minimum allowed coil-surface distance + block_size: Block size for memory efficiency. If None, uses full vmap (no chunking) + nfp: Number of field periods + stellsym: Whether stellarator symmetry is present Returns: - float: The curve-surface distance penalty value. + Scalar loss (sum over all relevant coil-surface pairs) """ - dists = jnp.sqrt(jnp.sum( - (gammac[:, None, :] - gammas[None, :, :])**2, axis=2)) - integralweight = jnp.linalg.norm(lc, axis=1)[:, None] \ - * jnp.linalg.norm(ns, axis=1)[None, :] - return jnp.mean(integralweight * jnp.maximum(minimum_distance-dists, 0)**2) - - - -#This is thr quickest way to get coil-coil distance (but I guess not the most efficient way for large sizes). -# In that case we would do the candidates method from simsopt entirely -def loss_linking_mnumber(x,dofs_curves,currents_scale,nfp,n_segments=60,stellsym=True,downsample=1): - coils=coils_from_dofs(x,dofs_curves, currents_scale, nfp,n_segments, stellsym) - #Since the quadpoints are the same for every curve then we can calculate the increment is constant for every curve - # (needs change if quadpoints are allowed to be different) - dphi=coils.quadpoints[1]-coils.quadpoints[0] - result=jnp.sum(jnp.triu(jax.vmap(jax.vmap(linking_number_pure,in_axes=(0,0,None,None,None)), - in_axes=(None,None,0,0,None))(coils.gamma[:,0:-1:downsample,:], - coils.gamma_dash[:,0:-1:downsample,:], - coils.gamma[:,0:-1:downsample,:], - coils.gamma_dash[:,0:-1:downsample,:], - dphi),k=1)) - return result - - -#This is thr quickest way to get coil-coil distance (but I guess not the most efficient way for large sizes). -# In that case we would do the candidates method from simsopt entirely -def loss_linking_mnumber_constarint(x,dofs_curves,currents_scale,nfp,n_segments=60,stellsym=True,downsample=1): - coils=coils_from_dofs(x,dofs_curves, currents_scale, nfp,n_segments, stellsym) - #Since the quadpoints are the same for every curve then we can calculate the increment is constant for every curve - # (needs change if quadpoints are allowed to be different) - dphi=coils.quadpoints[1]-coils.quadpoints[0] - result=jnp.triu(jax.vmap(jax.vmap(linking_number_pure,in_axes=(0,0,None,None,None)), - in_axes=(None,None,0,0,None))(coils.gamma[:,0:-1:downsample,:], - coils.gamma_dash[:,0:-1:downsample,:], - coils.gamma[:,0:-1:downsample,:], - coils.gamma_dash[:,0:-1:downsample,:], - dphi)+1.e-18,k=1) - #The 1.e-18 above is just to get all the correct values in the following mask - return result[result != 0.0].flatten() - -def linking_number_pure(gamma1, lc1, gamma2, lc2,dphi): - linking_number_ij=jnp.sum(jnp.abs(jax.vmap(integrand_linking_number, in_axes=(0, 0, 0, 0,None,None))(gamma1, lc1, gamma2, lc2,dphi,dphi)/ (4*jnp.pi))) - return linking_number_ij - -def integrand_linking_number(r1,dr1,r2,dr2,dphi1,dphi2): + n_coils = coils.gamma.shape[0] + n_points_coil = coils.gamma.shape[1] + surface_points = surface.gamma.reshape(-1, 3) + n_points_surface = surface_points.shape[0] + + # Only check unique coils for symmetry + if coils.stellsym: + n_unique_coils = n_coils // (2 * coils.nfp) + else: + n_unique_coils = n_coils // coils.nfp + n_unique_coils = max(1, n_unique_coils) + unique_coil_indices = jnp.arange(n_unique_coils) + + def single_coil_loss(idx): + gamma_i = coils.gamma[idx] + gamma_dash_i = coils.gamma_dash[idx] + gamma_dash_norm = jnp.linalg.norm(gamma_dash_i, axis=1) + n_points = gamma_i.shape[0] + + # If block_size is None, use full vmap (no chunking) + use_block_size = min(n_points_surface, n_points_surface if block_size is None else block_size) + n_blocks = (n_points_surface + use_block_size - 1) // use_block_size + padded_points = n_blocks * use_block_size + pad_width = padded_points - n_points_surface + + surface_point_blocks = jnp.pad(surface_points, ((0, pad_width), (0, 0))).reshape(n_blocks, use_block_size, 3) + valid_blocks = (jnp.arange(padded_points) < n_points_surface).reshape(n_blocks, use_block_size) + + def block_sum(block_surface_points, block_valid): + dists_block = jnp.linalg.norm(gamma_i[:, None, :] - block_surface_points[None, :, :], axis=2) + penalty_block = jnp.maximum(0, min_distance - dists_block) + weighted_penalty = jnp.square(penalty_block) * gamma_dash_norm[:, None] * block_valid[None, :] + return jnp.sum(weighted_penalty) + + total = jnp.sum(jax.vmap(block_sum)(surface_point_blocks, valid_blocks)) + norm = gamma_i.shape[0] * n_points_surface + return total / norm + + losses = jax.vmap(single_coil_loss)(unique_coil_indices) + return jnp.sum(losses) + + +# Blockwise vmap linking number loss (memory efficient, fully differentiable) +@partial(jit, static_argnames=["block_size"]) +def loss_linkingnumber(coils, candidates=None, block_size=None): + if candidates is None: + candidates = jnp.triu_indices(len(coils), k=1) + dphi = coils.curves.quadpoints[1] - coils.curves.quadpoints[0] + + def pair_linking(i, j): + gamma_i = coils.gamma[i] + gamma_dash_i = coils.gamma_dash[i] + gamma_j = coils.gamma[j] + gamma_dash_j = coils.gamma_dash[j] + n_points = gamma_j.shape[0] + + # If block_size is None, use full vmap (no chunking) + use_block_size = min(n_points, n_points if block_size is None else block_size) + n_blocks = (n_points + use_block_size - 1) // use_block_size + padded_points = n_blocks * use_block_size + pad_width = padded_points - n_points + + gamma_j_blocks = jnp.pad(gamma_j, ((0, pad_width), (0, 0))).reshape(n_blocks, use_block_size, 3) + gamma_dash_j_blocks = jnp.pad(gamma_dash_j, ((0, pad_width), (0, 0))).reshape(n_blocks, use_block_size, 3) + valid_blocks = (jnp.arange(padded_points) < n_points).reshape(n_blocks, use_block_size) + + def block_sum(block_gamma_j, block_gamma_dash_j, block_valid): + def integrand(r2, dr2): + diff = gamma_i - r2 + cross = jnp.cross(gamma_dash_i, dr2) + norm = jnp.linalg.norm(diff, axis=1) + return jnp.sum(diff * cross, axis=1) / (norm**3 + 1e-12) + + block_vals = jax.vmap(integrand, in_axes=(0, 0))(block_gamma_j, block_gamma_dash_j) + return jnp.sum(block_vals * block_valid[:, None]) + + total = jnp.sum(jax.vmap(block_sum)(gamma_j_blocks, gamma_dash_j_blocks, valid_blocks)) + linking = total * (dphi ** 2) / (4 * jnp.pi) + return jnp.abs(linking) + + losses = jax.vmap(pair_linking)(*candidates) + return jnp.sum(losses) + + + + +# Lorentz force loss: accepts Coils object, keyword args, JAX-friendly +@partial(jit, static_argnames=["p", "threshold", "block_size"]) +def loss_lorentz_force_coils(coils, p=1, threshold=0.5e6, block_size=None): """ - Compute the integrand for the linking number between two curves. - + Loss function penalizing Lorentz force on coils using Landreman-Hurwitz method. Args: - r1 (array-like): Points along the first curve. - dr1 (array-like): Tangent vectors along the first curve. - r2 (array-like): Points along the second curve. - dr2 (array-like): Tangent vectors along the second curve. - dphi1 (array-like): increments of quadpoints 1 - dphi2 (array-like): increments of quadpoints 2 - + coils: Coils object (with gamma, gamma_dash, gamma_dashdash, currents, quadpoints) + p: Power for penalty (default 1) + threshold: Force threshold (default 0.5e6) + block_size: Block size for memory efficiency. If None, uses full vmap (no chunking) Returns: - float: The integrand value for the linking number. + Scalar loss (sum over all coils) """ - return jnp.dot((r1-r2), jnp.cross(dr1, dr2)) / jnp.linalg.norm(r1-r2)**3*dphi1*dphi2 - - - -#Loss function penalizing force on coils using Landremann-Hurwitz method -def loss_lorentz_force_coils(x,dofs_curves,currents_scale,nfp,n_segments=60,stellsym=True,p=1,threshold=0.5e+6): - coils=coils_from_dofs(x,dofs_curves,currents_scale,nfp,n_segments, stellsym) - curves_indeces=jnp.arange(coils.gamma.shape[0]) - #We want to calculate tangeng cross [B_self + B_mutual] for each coil - #B_self is the self-field of the coil, B_mutual is the field from the other coils - force_penalty=jax.vmap(lp_force_pure,in_axes=(0,None,None,None,None,None,None,None))(curves_indeces,coils.gamma, - coils.gamma_dash,coils.gamma_dashdash,coils.currents,coils.quadpoints,p, threshold) - return force_penalty - + n_coils = coils.gamma.shape[0] + indices = jnp.arange(n_coils) + other_indices = jnp.array([ + [j for j in range(n_coils) if j != i] + for i in range(n_coils) + ], dtype=jnp.int32) + + def single_coil_loss(idx): + gamma_i = coils.gamma[idx] + gamma_dash_i = coils.gamma_dash[idx] + gamma_dashdash_i = coils.gamma_dashdash[idx] + current_i = coils.currents[idx] + quadpoints = coils.curves.quadpoints + curvature = Curves.compute_curvature(gamma_dash_i, gamma_dashdash_i) + regularization = regularization_circ(1. / jnp.mean(curvature)) + other_idx = other_indices[idx] + gamma_others = coils.gamma[other_idx] + gamma_dash_others = coils.gamma_dash[other_idx] + gamma_dashdash_others = coils.gamma_dashdash[other_idx] + currents_others = coils.currents[other_idx] + biot_savart = BiotSavart_from_gamma(gamma_others, gamma_dash_others, gamma_dashdash_others, currents_others) + n_points = gamma_i.shape[0] + use_block_size = min(n_points, n_points if block_size is None else block_size) + n_blocks = (n_points + use_block_size - 1) // use_block_size + padded_points = n_blocks * use_block_size + pad_width = padded_points - n_points + + gamma_i_blocks = jnp.pad(gamma_i, ((0, pad_width), (0, 0))).reshape(n_blocks, use_block_size, 3) + valid_blocks = (jnp.arange(padded_points) < n_points).reshape(n_blocks, use_block_size) + + def block_field(block_gamma_i, block_valid): + block_B = jax.vmap(biot_savart.B)(block_gamma_i) + return block_B * block_valid[:, None] + + block_B_mutual = jax.vmap(block_field)(gamma_i_blocks, valid_blocks).reshape(padded_points, 3)[:n_points] + block_gammadash_norm = jnp.linalg.norm(gamma_dash_i, axis=1) + block_tangent = gamma_dash_i / block_gammadash_norm[:, None] + block_B_self = B_regularized_pure( + gamma_i, gamma_dash_i, gamma_dashdash_i, + quadpoints, current_i, regularization + ) + block_force = jnp.cross(current_i * block_tangent, block_B_self + block_B_mutual) + block_force_norm = jnp.linalg.norm(block_force, axis=1) + total_penalty = jnp.sum(jnp.maximum(block_force_norm - threshold, 0) ** p * block_gammadash_norm) + return total_penalty * (1. / p) + + penalties = jax.vmap(single_coil_loss)(indices) + return jnp.sum(penalties) -def lp_force_pure(index,gamma, gamma_dash,gamma_dashdash,currents,quadpoints,p, threshold): - """Pure function for minimizing the Lorentz force on a coil. - """ - regularization = regularization_circ(1./jnp.average(compute_curvature( gamma_dash.at[index].get(), gamma_dashdash.at[index].get()))) - B_mutual=jax.vmap(BiotSavart_from_gamma(jnp.roll(gamma, -index, axis=0)[1:], - jnp.roll(gamma_dash, -index, axis=0)[1:], - jnp.roll(gamma_dashdash, -index, axis=0)[1:], - jnp.roll(currents, -index, axis=0)[1:]).B,in_axes=0)(gamma[index]) - B_self = B_regularized_pure(gamma.at[index].get(),gamma_dash.at[index].get(), gamma_dashdash.at[index].get(), quadpoints, currents[index], regularization) - gammadash_norm = jnp.linalg.norm(gamma_dash.at[index].get(), axis=1)[:, None] - tangent = gamma_dash.at[index].get() / gammadash_norm - force = jnp.cross(currents.at[index].get() * tangent, B_self + B_mutual) - force_norm = jnp.linalg.norm(force, axis=1)[:, None] - return (jnp.sum(jnp.maximum(force_norm - threshold, 0)**p * gammadash_norm))*(1./p) - - def B_regularized_singularity_term(rc_prime, rc_prime_prime, regularization): """The term in the regularized Biot-Savart law in which the near-singularity @@ -663,4 +541,8 @@ def rectangular_xsection_delta(a, b): # # bdotn_over_b_loss = jnp.sum(jnp.abs(bdotn_over_b)) -# return bdotn_over_b_loss \ No newline at end of file +# return bdotn_over_b_loss + + + +####################### SURFACE GEOMETRY LOSSES ########################## diff --git a/essos/optimization.py b/essos/optimization.py index fb1a24bb..21a29059 100644 --- a/essos/optimization.py +++ b/essos/optimization.py @@ -5,7 +5,7 @@ from functools import partial from essos.coils import Curves, Coils from scipy.optimize import least_squares, minimize -from essos.fields import near_axis +from pyqsc_jax.near_axis import near_axis from essos.surfaces import SurfaceRZFourier def new_nearaxis_from_x_and_old_nearaxis(new_field_nearaxis_x, field_nearaxis): @@ -64,7 +64,7 @@ def optimize_loss_function(func, initial_dofs, coils, tolerance_optimization=1e- dofs_currents = result.x[len_dofs_curves:-len(surface_all.x)] curves = Curves(dofs_curves, n_segments, nfp, stellsym) new_coils = Coils(curves=curves, currents=dofs_currents * coils.currents_scale) - new_surface = SurfaceRZFourier(rc=surface_all.rc, zs=surface_all.zs, nfp=nfp, range_torus=surface_all.range_torus, nphi=surface_all.nphi, ntheta=surface_all.ntheta) + new_surface = SurfaceRZFourier(rc=surface_all.rc, zs=surface_all.zs, nfp=nfp, range_torus=surface_all.range_torus, nphi=surface_all.nphi, ntheta=surface_all.ntheta,mpol=surface_all.mpol,ntor=surface_all.ntor) new_surface.dofs = result.x[-len(surface_all.x):] return new_coils, new_surface elif 'surface_all' in kwargs and 'field_nearaxis' in kwargs and len(initial_dofs) == len(coils.x) + len(kwargs['surface_all'].x) + len(kwargs['field_nearaxis'].x): @@ -73,7 +73,7 @@ def optimize_loss_function(func, initial_dofs, coils, tolerance_optimization=1e- dofs_currents = result.x[len_dofs_curves:-len(surface_all.x)-len(field_nearaxis.x)] curves = Curves(dofs_curves, n_segments, nfp, stellsym) new_coils = Coils(curves=curves, currents=dofs_currents * coils.currents_scale) - new_surface = SurfaceRZFourier(rc=surface_all.rc, zs=surface_all.zs, nfp=nfp, range_torus=surface_all.range_torus, nphi=surface_all.nphi, ntheta=surface_all.ntheta) + new_surface = SurfaceRZFourier(rc=surface_all.rc, zs=surface_all.zs, nfp=nfp, range_torus=surface_all.range_torus, nphi=surface_all.nphi, ntheta=surface_all.ntheta,mpol=surface_all.mpol,ntor=surface_all.ntor) new_surface.dofs = result.x[-len(surface_all.x)-len(field_nearaxis.x):-len(field_nearaxis.x)] new_field_nearaxis = new_nearaxis_from_x_and_old_nearaxis(result.x[-len(field_nearaxis.x):], field_nearaxis) return new_coils, new_surface, new_field_nearaxis diff --git a/essos/surfaces.py b/essos/surfaces.py index 0048e3cf..3dd31460 100644 --- a/essos/surfaces.py +++ b/essos/surfaces.py @@ -1,36 +1,92 @@ from functools import partial +import jax import jax.numpy as jnp from jax.scipy.interpolate import RegularGridInterpolator -from jax import jit, vmap, devices, device_put +from jax import tree_util, jit, vmap, devices, device_put from jax.sharding import Mesh, NamedSharding, PartitionSpec +from jax.experimental.pjit import pjit from essos.plot import fix_matplotlib_3d import jaxkd mesh = Mesh(devices(), ("dev",)) -sharding = NamedSharding(mesh, PartitionSpec("dev", None)) +sharding = NamedSharding(mesh, PartitionSpec("dev")) -@partial(jit, static_argnames=['surface','field']) + +@jit +def toroidal_flux(surface, field, idx=0) -> jnp.ndarray: + curve = surface.gamma[idx] + dl = jnp.roll(curve, -1, axis=0) - curve + A_vals = vmap(field.A)(curve) + Adl = jnp.sum(A_vals * dl, axis=1) + tf = jnp.sum(Adl) + #curve = surface.gamma[idx] + #dl = surface.gammadash_theta[idx] + #A_vals = vmap(field.A)(curve) + #Adl = jnp.sum(A_vals * dl, axis=1)/surface.ntheta + #tf = jnp.sum(Adl) + return tf + +@jit +def poloidal_flux(surface, field, idx=0) -> jnp.ndarray: + curve = surface.gamma[:,idx,:] + dl = jnp.roll(curve, -1, axis=0) - curve + A_vals = vmap(field.A)(curve) + Adl = jnp.sum(A_vals * dl, axis=1) + tf = jnp.sum(Adl) + #curve = surface.gamma[:,idx,:] + #dl = surface.gammadash_phi[:,idx,:] + #A_vals = vmap(field.A)(curve) + #Adl = jnp.sum(A_vals * dl, axis=1)/surface.nphi + #tf = jnp.sum(Adl) + return tf + +# @jit +@partial(pjit, in_shardings=(sharding, None), out_shardings=sharding) def B_on_surface(surface, field): ntheta = surface.ntheta nphi = surface.nphi gamma = surface.gamma gamma_reshaped = gamma.reshape(nphi * ntheta, 3) - gamma_sharded = device_put(gamma_reshaped, sharding) - B_on_surface = jit(vmap(field.B), in_shardings=sharding, out_shardings=sharding)(gamma_sharded) - B_on_surface = B_on_surface.reshape(nphi, ntheta, 3) - return B_on_surface -@partial(jit, static_argnames=['surface','field']) + # Map field.B over all positions + B_on_surface = vmap(field.B)(gamma_reshaped) + + return B_on_surface.reshape(nphi, ntheta, 3) + + +@jit def BdotN(surface, field): B_surface = B_on_surface(surface, field) B_dot_n = jnp.sum(B_surface * surface.unitnormal, axis=2) return B_dot_n -@partial(jit, static_argnames=['surface','field']) -def BdotN_over_B(surface, field): - B_surface = B_on_surface(surface, field) - B_dot_n = jnp.sum(B_surface * surface.unitnormal, axis=2) - return B_dot_n / jnp.linalg.norm(B_surface, axis=2) +@jit +def BdotN_over_B(surface, field, **kwargs): + return BdotN(surface, field) / jnp.linalg.norm(B_on_surface(surface, field), axis=2) + +@jit +def _squared_flux_local(surface, field): + return 0.5 * jnp.mean(BdotN(surface, field)**2 / jnp.sum(B_on_surface(surface, field)**2, axis=2) + * surface.area_element) + +@jit +def _squared_flux_global(surface, field): + return 0.5 * jnp.mean(BdotN(surface, field)**2 * surface.area_element) + +@jit +def _squared_flux_normalized(surface, field): + return 0.5 * jnp.mean(BdotN(surface, field)**2 * surface.area_element) / \ + jnp.mean(jnp.sum(B_on_surface(surface, field)**2, axis=2) * surface.area_element) + +def SquaredFlux(surface, field, definition='local'): + if definition == 'local': + return _squared_flux_local(surface, field) + elif definition == 'quadratic flux': + return _squared_flux_global(surface, field) + elif definition == 'normalized': + return _squared_flux_normalized(surface, field) + else: + raise ValueError(f"Unknown definition: {definition}") def nested_lists_to_array(ll): """ @@ -47,187 +103,420 @@ def nested_lists_to_array(ll): for jm, l in enumerate(ll): arr = arr.at[jm, :len(l)].set(jnp.array([x if x is not None else 0 for x in l])) return arr + class SurfaceRZFourier: - def __init__(self, vmec=None, s=1, ntheta=30, nphi=30, close=True, range_torus='full torus', - rc=None, zs=None, nfp=None): - if rc is not None: - self.rc = rc - self.zs = zs - self.nfp = nfp - self.mpol = rc.shape[0] - self.ntor = (rc.shape[1] - 1) // 2 - m1d = jnp.arange(self.mpol) - n1d = jnp.arange(-self.ntor, self.ntor + 1) - n2d, m2d = jnp.meshgrid(n1d, m1d) - self.xm = m2d.flatten()[self.ntor:] - self.xn = self.nfp*n2d.flatten()[self.ntor:] - indices = jnp.array([self.xm, self.xn / self.nfp + self.ntor], dtype=int).T - self.rmnc_interp = self.rc[indices[:, 0], indices[:, 1]] - self.zmns_interp = self.zs[indices[:, 0], indices[:, 1]] - elif isinstance(vmec, str): - self.input_filename = vmec - import f90nml - all_namelists = f90nml.read(vmec) - nml = all_namelists['indata'] - if 'nfp' in nml: - self.nfp = nml['nfp'] - else: - self.nfp = 1 - rc = nested_lists_to_array(nml['rbc']) - zs = nested_lists_to_array(nml['zbs']) - rbc_first_n = nml.start_index['rbc'][0] - rbc_last_n = rbc_first_n + rc.shape[1] - 1 - zbs_first_n = nml.start_index['zbs'][0] - zbs_last_n = zbs_first_n + zs.shape[1] - 1 - self.ntor = jnp.max(jnp.abs(jnp.array([rbc_first_n, rbc_last_n, zbs_first_n, zbs_last_n], dtype='i'))) - rbc_first_m = nml.start_index['rbc'][1] - rbc_last_m = rbc_first_m + rc.shape[0] - 1 - zbs_first_m = nml.start_index['zbs'][1] - zbs_last_m = zbs_first_m + zs.shape[0] - 1 - self.mpol = max(rbc_last_m, zbs_last_m) - self.rc = jnp.zeros((self.mpol, 2 * self.ntor + 1)) - self.zs = jnp.zeros((self.mpol, 2 * self.ntor + 1)) - m_indices_rc = jnp.arange(rc.shape[0]) + nml.start_index['rbc'][1] - n_indices_rc = jnp.arange(rc.shape[1]) + nml.start_index['rbc'][0] + self.ntor - self.rc = self.rc.at[m_indices_rc[:, None], n_indices_rc].set(rc) - m_indices_zs = jnp.arange(zs.shape[0]) + nml.start_index['zbs'][1] - n_indices_zs = jnp.arange(zs.shape[1]) + nml.start_index['zbs'][0] + self.ntor - self.zs = self.zs.at[m_indices_zs[:, None], n_indices_zs].set(zs) - m1d = jnp.arange(self.mpol) - n1d = jnp.arange(-self.ntor, self.ntor + 1) - n2d, m2d = jnp.meshgrid(n1d, m1d) - self.xm = m2d.flatten()[self.ntor:] - self.xn = self.nfp*n2d.flatten()[self.ntor:] - indices = jnp.array([self.xm, self.xn / self.nfp + self.ntor], dtype=int).T - self.rmnc_interp = self.rc[indices[:, 0], indices[:, 1]] - self.zmns_interp = self.zs[indices[:, 0], indices[:, 1]] - else: - try: - self.nfp = vmec.nfp - self.bmnc = vmec.bmnc - self.xm = vmec.xm - self.xn = vmec.xn - self.rmnc = vmec.rmnc - self.zmns = vmec.zmns - self.xm_nyq = vmec.xm_nyq - self.xn_nyq = vmec.xn_nyq - self.len_xm_nyq = len(self.xm_nyq) - self.ns = vmec.ns - self.s_full_grid = vmec.s_full_grid - self.ds = vmec.ds - self.s_half_grid = vmec.s_half_grid - self.r_axis = vmec.r_axis - self.rmnc_interp = vmap(lambda row: jnp.interp(s, self.s_full_grid, row, left='extrapolate'), in_axes=1)(self.rmnc) - self.zmns_interp = vmap(lambda row: jnp.interp(s, self.s_full_grid, row, left='extrapolate'), in_axes=1)(self.zmns) - self.bmnc_interp = vmap(lambda row: jnp.interp(s, self.s_half_grid, row, left='extrapolate'), in_axes=1)(self.bmnc[1:, :]) - self.mpol = vmec.mpol - self.ntor = vmec.ntor - self.num_dofs = 2 * (self.mpol + 1) * (2 * self.ntor + 1) - self.ntor - (self.ntor + 1) - shape = (int(jnp.max(self.xm)) + 1, int(jnp.max(self.xn)) + 1) - self.rc = jnp.zeros(shape) - self.zs = jnp.zeros(shape) - indices = jnp.array([self.xm, self.xn / self.nfp + self.ntor], dtype=int).T - self.rc = self.rc.at[indices[:, 0], indices[:, 1]].set(self.rmnc_interp) - self.zs = self.zs.at[indices[:, 0], indices[:, 1]].set(self.zmns_interp) - except: - raise ValueError("vmec must be a Vmec object or a string pointing to a VMEC input file.") - self.ntheta = ntheta - self.nphi = nphi - self.range_torus = range_torus - if range_torus == 'full torus': div = 1 - else: div = self.nfp - if range_torus == 'half period': end_val = 0.5 - else: end_val = 1.0 - self.quadpoints_theta = jnp.linspace(0, 2 * jnp.pi, num=self.ntheta, endpoint=True if close else False) - self.quadpoints_phi = jnp.linspace(0, 2 * jnp.pi * end_val / div, num=self.nphi, endpoint=True if close else False) - self.theta_2d, self.phi_2d = jnp.meshgrid(self.quadpoints_theta, self.quadpoints_phi) - self.num_dofs_rc = len(jnp.ravel(self.rc)[self.ntor:]) - self.num_dofs_zs = len(jnp.ravel(self.zs)[self.ntor:]) - self._dofs = jnp.concatenate((jnp.ravel(self.rc)[self.ntor:], jnp.ravel(self.zs)[self.ntor:])) + def __init__(self, rc, zs, nfp, mpol, ntor, ntheta=30, nphi=30, close=True, range_torus='full torus', + scaling_type=2, scaling_factor=0): + """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." + assert isinstance(ntor, int) and ntor >= 0, "ntor must be a non-negative integer." + assert isinstance(ntheta, int) and ntheta > 0, "ntheta must be a positive integer." + assert isinstance(nphi, int) and nphi > 0, "nphi must be a positive integer." + assert isinstance(close, bool), "close must be a boolean." + assert range_torus in ['full torus', 'half period'], f"Unknown range_torus: {range_torus}. Choose 'full torus' or 'half period'." + self._initialize_state( + rc, + zs, + nfp, + mpol, + ntor, + ntheta, + nphi, + close, + range_torus, + self._normalize_scaling_type(scaling_type), + scaling_factor, + ) + + def _initialize_state(self, rc, zs, nfp, mpol, ntor, ntheta, nphi, close, range_torus, scaling_type, scaling_factor): + self._rc = rc + self._zs = zs + self._nfp = nfp + self._mpol = mpol + self._ntor = ntor + + self._gamma = None + self._gammadash_theta = None + self._gammadash_phi = None + self._normal = None + self._unitnormal = None + self._area_element = None + self._xm = None + self._xn = None + + self._ntheta = ntheta + self._nphi = nphi + self._close = close + self._range_torus = range_torus + + self._quadpoints_theta = None + self._quadpoints_phi = None + self._theta2d = None + self._phi2d = None + self._angles = None + self._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." + ) + + @staticmethod + def _compute_scaling(xm, xn, scaling_type, scaling_factor): + return jnp.exp(scaling_factor * jnp.linalg.norm(jnp.vstack([xm, xn]), ord=scaling_type, axis=0)) + + + @classmethod + def from_input_file(cls, file, ntheta=30, nphi=30, close=True, range_torus='full torus'): + from f90nml import Parser + nml = Parser().read(file)['indata'] + + nfp = nml["nfp"] if "nfp" in nml else 1 + mpol = nml['mpol'] + ntor = nml['ntor'] - self.angles = jnp.einsum('i,jk->ijk', self.xm, self.theta_2d) - jnp.einsum('i,jk->ijk', self.xn, self.phi_2d) + rc = jnp.ravel(nested_lists_to_array(nml['rbc']))[2:] + zs = jnp.ravel(nested_lists_to_array(nml['zbs']))[2:] + + surface = cls(rc, zs, nfp, mpol, ntor, ntheta=ntheta, nphi=nphi, close=close, range_torus=range_torus) + return surface - (self._gamma, self._gammadash_theta, self._gammadash_phi, - self._normal, self._unitnormal) = self._set_gamma(self.rmnc_interp, self.zmns_interp) + @classmethod + def from_vmec(cls, vmec, s=1, ntheta=30, nphi=30, close=True, range_torus='full torus'): + nfp = vmec.nfp + mpol = vmec.mpol + ntor = vmec.ntor + + s_full_grid = vmec.s_full_grid + rc = vmap(lambda row: jnp.interp(s, s_full_grid, row, left='extrapolate'), in_axes=1)(vmec.rmnc) + zs = vmap(lambda row: jnp.interp(s, s_full_grid, row, left='extrapolate'), in_axes=1)(vmec.zmns) + + surface = cls(rc, zs, nfp, mpol, ntor, ntheta=ntheta, nphi=nphi, close=close, range_torus=range_torus) + surface._xm = vmec.xm + surface._xn = vmec.xn + + return surface + + @classmethod + def from_wout_file(cls, file, s=1, ntheta=30, nphi=30, close=True, range_torus='full torus'): + from netCDF4 import Dataset + nc = Dataset(file) + + nfp = int(nc.variables["nfp"][0]) + xm = jnp.array(nc.variables["xm"][:]) + xn = jnp.array(nc.variables["xn"][:]) + mpol = int(jnp.max(xm)) + ntor = int(jnp.max(jnp.abs(xn)) / nfp) - if hasattr(self, 'bmnc'): - self._AbsB = self._set_AbsB() + ns = nc.variables["ns"][0] + s_full_grid = jnp.linspace(0, 1, ns) + rc = vmap(lambda row: jnp.interp(s, s_full_grid, row, left='extrapolate'), in_axes=1)(jnp.array(nc.variables["rmnc"][:])) + zs = vmap(lambda row: jnp.interp(s, s_full_grid, row, left='extrapolate'), in_axes=1)(jnp.array(nc.variables["zmns"][:])) + + surface = cls(rc, zs, nfp, mpol, ntor, ntheta=ntheta, nphi=nphi, close=close, range_torus=range_torus) + surface._xm = xm + surface._xn = xn + + return surface + + # reset_cache method + def reset_cache(self): + self._gamma = None + self._gammadash_theta = None + self._gammadash_phi = None + self._normal = None + self._unitnormal = None + self._area_element = None + self._xm = None + self._xn = None + self._angles = None + + # reset_mesh method + def reset_mesh(self): + self._quadpoints_theta = None + self._quadpoints_phi = None + self._theta2d = None + self._phi2d = None + self._angles = None + + # rc property and setter + @property + def rc(self): + return self._rc + + @rc.setter + def rc(self, new_rc): + self._rc = new_rc + self.reset_cache() + + # zs property and setter + @property + def zs(self): + return self._zs + @zs.setter + def zs(self, new_zs): + self._zs = new_zs + self.reset_cache() + + # nfp property + @property + def nfp(self): + return self._nfp + + # mpol property + @property + def mpol(self): + return self._mpol + + # ntor property + @property + def ntor(self): + return self._ntor + + # xm property + @property + def xm(self): + if self._xm is None: + self._xm = jnp.repeat(jnp.arange(self.mpol + 1), 2 * self.ntor + 1)[self.ntor:] + return self._xm + + # xn property + @property + def xn(self): + if self._xn is None: + self._xn = self.nfp * jnp.tile(jnp.arange(-self.ntor, self.ntor + 1), self.mpol + 1)[self.ntor:] + return self._xn + + # _ntheta property and setter + @property + def ntheta(self): + return self._ntheta + + @ntheta.setter + def ntheta(self, new_ntheta): + self._ntheta = new_ntheta + self.reset_mesh() + + # n_phi property and setter + @property + def nphi(self): + return self._nphi + + @nphi.setter + def nphi(self, new_nphi): + self._nphi = new_nphi + self.reset_mesh() + + # close property and setter + @property + def close(self): + return self._close + + @close.setter + def close(self, new_close): + self._close = new_close + self.reset_mesh() + + # range_torus property and setter + @property + def range_torus(self): + return self._range_torus + + @range_torus.setter + def range_torus(self, new_range): + self._range_torus = new_range + self.reset_mesh() + + # _compute_meshgrid method + @jit + def _compute_meshgrid(self): + if self.range_torus == "full torus": + div, end_val = 1., 1. + elif self.range_torus == "half period": + div, end_val = self.nfp, 0.5 + quadpoints_theta = jnp.linspace(0, 2 * jnp.pi, num=self.ntheta, endpoint=self.close) + quadpoints_phi = jnp.linspace(0, 2 * jnp.pi * end_val / div, num=self.nphi, endpoint=self.close) + theta2d, phi2d = jnp.meshgrid(quadpoints_theta, quadpoints_phi) + return quadpoints_theta, quadpoints_phi, theta2d, phi2d + + # theta2d property + @property + def theta2d(self): + if self._theta2d is None: + self._quadpoints_theta, self._quadpoints_phi, self._theta2d, self._phi2d = self._compute_meshgrid() + return self._theta2d + + # phi2d property + @property + def phi2d(self): + if self._phi2d is None: + self._quadpoints_theta, self._quadpoints_phi, self._theta2d, self._phi2d = self._compute_meshgrid() + return self._phi2d + + # angles property + @property + def angles(self): + if self._angles is None: + self._angles = jnp.einsum('i,jk->ijk', self.xm, self.theta2d) - jnp.einsum('i,jk->ijk', self.xn, self.phi2d) + return self._angles + + # 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 + + # scaling property + @property + def scaling(self): + """Mode-by-mode scaling ``exp(scaling_factor * ||(xm, xn)||)``.""" + if self._scaling is None: + scaling = self._compute_scaling(self.xm, self.xn, self.scaling_type, self.scaling_factor) + if not isinstance(scaling, jax.core.Tracer): + self._scaling = scaling + return scaling + return self._scaling + + # dofs property and setter @property def dofs(self): - return self._dofs + return jnp.hstack([self.rc * self.scaling, self.zs * self.scaling]) @dofs.setter def dofs(self, new_dofs): - self._dofs = new_dofs - self.rc = jnp.concatenate((jnp.zeros(self.ntor),new_dofs[:self.num_dofs_rc])).reshape(self.rc.shape) - self.zs = jnp.concatenate((jnp.zeros(self.ntor),new_dofs[self.num_dofs_rc:])).reshape(self.zs.shape) - indices = jnp.array([self.xm, self.xn / self.nfp + self.ntor], dtype=int).T - self.rmnc_interp = self.rc[indices[:, 0], indices[:, 1]] - self.zmns_interp = self.zs[indices[:, 0], indices[:, 1]] - (self._gamma, self._gammadash_theta, self._gammadash_phi, - self._normal, self._unitnormal) = self._set_gamma(self.rmnc_interp, self.zmns_interp) - # if hasattr(self, 'bmnc'): - # self._AbsB = self._set_AbsB() + self._rc = new_dofs[:self.rc.size] / self.scaling + self._zs = new_dofs[self.rc.size:] / self.scaling + self.reset_cache() - @partial(jit, static_argnames=['self']) - def _set_gamma(self, rmnc_interp, zmns_interp): - phi_2d = self.phi_2d + # _compute_gamma method + @jit + def _compute_gamma(self): angles = self.angles - sin_angles = jnp.sin(angles) cos_angles = jnp.cos(angles) - r_coordinate = jnp.einsum('i,ijk->jk', rmnc_interp, cos_angles) - z_coordinate = jnp.einsum('i,ijk->jk', zmns_interp, sin_angles) - gamma = jnp.transpose(jnp.array([r_coordinate * jnp.cos(phi_2d), r_coordinate * jnp.sin(phi_2d), z_coordinate]), (1, 2, 0)) + phi2d = self.phi2d + sin_phi2d = jnp.sin(phi2d) + cos_phi2d = jnp.cos(phi2d) + rc = self.rc; zs = self.zs; xm = self.xm; xn = self.xn - dX_dtheta = jnp.einsum('i,ijk,i->jk', -self.xm, sin_angles, rmnc_interp) * jnp.cos(phi_2d) - dY_dtheta = jnp.einsum('i,ijk,i->jk', -self.xm, sin_angles, rmnc_interp) * jnp.sin(phi_2d) - dZ_dtheta = jnp.einsum('i,ijk,i->jk', self.xm, cos_angles, zmns_interp) - gammadash_theta = 2*jnp.pi*jnp.transpose(jnp.array([dX_dtheta, dY_dtheta, dZ_dtheta]), (1, 2, 0)) + R = jnp.einsum('i,ijk->jk', rc, cos_angles) + Z = jnp.einsum('i,ijk->jk', zs, sin_angles) + X = R * cos_phi2d + Y = R * sin_phi2d + gamma = jnp.stack([X, Y, Z], axis=-1) - dX_dphi = jnp.einsum('i,ijk,i->jk', self.xn, sin_angles, rmnc_interp) * jnp.cos(phi_2d) - r_coordinate * jnp.sin(phi_2d) - dY_dphi = jnp.einsum('i,ijk,i->jk', self.xn, sin_angles, rmnc_interp) * jnp.sin(phi_2d) + r_coordinate * jnp.cos(phi_2d) - dZ_dphi = jnp.einsum('i,ijk,i->jk', -self.xn, cos_angles, zmns_interp) - gammadash_phi = 2*jnp.pi*jnp.transpose(jnp.array([dX_dphi, dY_dphi, dZ_dphi]), (1, 2, 0)) + dR_dtheta = -jnp.einsum('i,ijk->jk', xm * rc, sin_angles) + dZ_dtheta = jnp.einsum('i,ijk->jk', xm * zs, cos_angles) + dX_dtheta = dR_dtheta * cos_phi2d + dY_dtheta = dR_dtheta * sin_phi2d + gammadash_theta = jnp.stack([dX_dtheta, dY_dtheta, dZ_dtheta], axis=-1) - normal = jnp.cross(gammadash_phi, gammadash_theta, axis=2) - unitnormal = normal / jnp.linalg.norm(normal, axis=2, keepdims=True) + dR_dphi = jnp.einsum('i,ijk->jk', xn*rc, sin_angles) + dZ_dphi = -jnp.einsum('i,ijk->jk', xn*zs, cos_angles) + dX_dphi = dR_dphi * cos_phi2d - R * sin_phi2d + dY_dphi = dR_dphi * sin_phi2d + R * cos_phi2d + gammadash_phi = jnp.stack([dX_dphi, dY_dphi, dZ_dphi], axis=-1) - return (gamma, gammadash_theta, gammadash_phi, normal, unitnormal) - - @partial(jit, static_argnames=['self']) - def _set_AbsB(self): - angles_nyq = jnp.einsum('i,jk->ijk', self.xm_nyq, self.theta_2d) - jnp.einsum('i,jk->ijk', self.xn_nyq, self.phi_2d) - AbsB = jnp.einsum('i,ijk->jk', self.bmnc_interp, jnp.cos(angles_nyq)) - return AbsB + return gamma, gammadash_theta, gammadash_phi + # gamma, gammadash_theta, gammadash_phi properties @property def gamma(self): + if self._gamma is None: + self._gamma, self._gammadash_theta, self._gammadash_phi = self._compute_gamma() return self._gamma @property def gammadash_theta(self): + if self._gammadash_theta is None: + self._gamma, self._gammadash_theta, self._gammadash_phi = self._compute_gamma() return self._gammadash_theta @property def gammadash_phi(self): + if self._gammadash_phi is None: + self._gamma, self._gammadash_theta, self._gammadash_phi = self._compute_gamma() return self._gammadash_phi + + # _compute_properties method + @jit + def _compute_properties(self): + normal = jnp.cross(self.gammadash_theta, self.gammadash_phi, axis=2) + unitnormal = normal / jnp.linalg.norm(normal, axis=2, keepdims=True) + area_element = jnp.linalg.norm(normal, axis=2) + return normal, unitnormal, area_element + # normal, unitnormal, area_element properties @property def normal(self): + if self._normal is None: + self._normal, self._unitnormal, self._area_element = self._compute_properties() return self._normal @property def unitnormal(self): + if self._unitnormal is None: + self._normal, self._unitnormal, self._area_element = self._compute_properties() return self._unitnormal @property - def AbsB(self): - return self._AbsB - + def area_element(self): + if self._area_element is None: + self._normal, self._unitnormal, self._area_element = self._compute_properties() + return self._area_element + + # TODO: remove x property. This is a placeholder for compatibility with the examples that need to be updated. + # x property and setter @property def x(self): return self.dofs @@ -235,7 +524,101 @@ def x(self): @x.setter def x(self, new_dofs): self.dofs = new_dofs - + + @property + def volume(self): + + xyz = self.gamma # shape: (nphi, ntheta, 3) + n = self.normal # shape: (nphi, ntheta, 3) + + integrand = jnp.sum(xyz * n, axis=2) # dot(x, n), shape: (nphi, ntheta) + volume = jnp.mean(integrand) / 3.0 + return volume + + @property + def area(self): + #n = self.normal # (nphi, ntheta, 3) + #norm_n = jnp.linalg.norm(n, axis=2) # shape: (nphi, ntheta) + #avg_area = jnp.mean(norm_n) + #return avg_area + n = self.normal # shape: (nphi, ntheta, 3) + norm_n = jnp.linalg.norm(n, axis=2) + + dphi = 2 * jnp.pi / self.nphi + dtheta = 2 * jnp.pi / self.ntheta + + area = jnp.sum(norm_n) * dphi * dtheta + return area + + # def change_resolution(self, mpol: int, ntor: int, ntheta=None, nphi=None,close=True): + # """ + # Change the values of `mpol` and `ntor`. + # New Fourier coefficients are zero by default. + # Old coefficients outside the new range are discarded. + # """ + # rc_old, zs_old = self.rc, self.zs + # mpol_old, ntor_old = self.mpol, self.ntor + # if ntheta is not None: + # self.ntheta = ntheta + # else: + # ntheta = self.ntheta + + # if nphi is not None: + # self.nphi = nphi + # else: + # nphi = self.nphi + + # #rc_new = jnp.zeros((mpol, 2 * ntor + 1)) + # #zs_new = jnp.zeros((mpol, 2 * ntor + 1)) + # rc_new = jnp.zeros(((mpol+1)*( 2 * ntor + 1)-ntor)) + # zs_new = jnp.zeros(((mpol+1)*( 2 * ntor + 1)-ntor)) + # m_keep = min(mpol_old, mpol) + # n_keep = min(ntor_old, ntor) + + # xm_old=self.xm + # xn_old=self.xn + # self.xm = jnp.repeat(jnp.arange(mpol+1), 2*ntor+1)[ntor:] + # self.xn = self.nfp*jnp.tile(jnp.arange(-ntor, ntor + 1), mpol+1)[ntor:] + # # Copy overlapping region + # for l in range(len(self.xm)): + # if self.xm[l]<=m_keep and jnp.abs(self.xn[l]/self.nfp)<=n_keep: + # index=self.xm[l]*(ntor_old*2+1)-self.xn[l]//self.nfp + # rc_new=rc_new.at[l].set(self.rc[index]) + # zs_new=zs_new.at[l].set(self.zs[index]) + + + # # Update attributes + # self.mpol, self.ntor = mpol, ntor + # self.rc, self.zs = rc_new, zs_new + + # self.rmnc_interp = self.rc + # self.zmns_interp = self.zs + + # # Update degrees of freedom + # self.num_dofs_rc = len(jnp.ravel(self.rc)) + # self.num_dofs_zs = len(jnp.ravel(self.zs)) + # self._dofs = jnp.concatenate((self.rescaling_function(jnp.ravel(self.rc)), self.rescaling_function(jnp.ravel(self.zs)))) + + # # Recompute angles and geometry + # if self.range_torus == 'full torus': div = 1 + # else: div = self.nfp + # if self.range_torus == 'half period': end_val = 0.5 + # else: end_val = 1.0 + # self.quadpoints_theta = jnp.linspace(0, 2 * jnp.pi, num=ntheta, endpoint=True if close else False) + # self.quadpoints_phi = jnp.linspace(0, 2 * jnp.pi * end_val / div, num=nphi, endpoint=True if close else False) + # self.theta_2d, self.phi_2d = jnp.meshgrid(self.quadpoints_theta, self.quadpoints_phi) + + # self.angles = (jnp.einsum('i,jk->ijk', self.xm, self.theta_2d)- jnp.einsum('i,jk->ijk', self.xn, self.phi_2d)) + # (self._gamma, self._gammadash_theta, self._gammadash_phi, + # self._normal, self._unitnormal) = self._set_gamma(self.rmnc_interp, self.zmns_interp) + + + # # Recompute AbsB if available + # if hasattr(self, 'bmnc'): + # self._AbsB = self._set_AbsB() + + # return self + def plot(self, ax=None, show=True, close=False, axis_equal=True, **kwargs): if close: raise NotImplementedError("Call close=True when instantiating the VMEC/SurfaceRZFourier object.") @@ -246,7 +629,7 @@ def plot(self, ax=None, show=True, close=False, axis_equal=True, **kwargs): if ax is None or ax.name != "3d": fig = plt.figure() ax = fig.add_subplot(projection='3d') - + boundary = self.gamma if hasattr(self, 'bmnc'): @@ -299,15 +682,11 @@ def to_vmec(self, filename): nml += 'LASYM = .FALSE.\n' nml += f'NFP = {self.nfp}\n' - for m in range(self.mpol + 1): - nmin = -self.ntor - if m == 0: - nmin = 0 - for n in range(nmin, self.ntor + 1): - rc = self.rc[m, n + self.ntor] - zs = self.zs[m, n + self.ntor] - if jnp.abs(rc) > 0 or jnp.abs(zs) > 0: - nml += f"RBC({n:4d},{m:4d}) ={rc:23.15e}, ZBS({n:4d},{m:4d}) ={zs:23.15e}\n" + # Copy overlapping region + for l in range(len(self.xm)): + rc = self.rc[l] + zs = self.zs[l] + nml += f"RBC({self.xn[l]:4d},{self.xm[l]:4d}) ={rc:23.15e}, ZBS({self.xn[l]:4d},{self.xm[l]:4d}) ={zs:23.15e}\n" nml += '/\n' with open(filename, 'w') as f: @@ -329,6 +708,63 @@ def mean_cross_sectional_area(self): mean_cross_sectional_area = jnp.abs(jnp.mean(jnp.sqrt(x2y2) * dZ_dtheta * detJ))/(2 * jnp.pi) return mean_cross_sectional_area + def _tree_flatten(self): + if hasattr(self._rc, "shape") and hasattr(self._zs, "shape"): + children = (self.rc * self.scaling, self.zs * self.scaling) # arrays / dynamic values + else: + children = (self._rc, self._zs) + aux_data = {"nfp": self._nfp, + "mpol": self._mpol, + "ntor": self._ntor, + "ntheta": self._ntheta, + "nphi": self._nphi, + "close": self._close, + "range_torus": self._range_torus, + "scaling_type": self._scaling_type, + "scaling_factor": self._scaling_factor} # static values + return (children, aux_data) + + @classmethod + def _tree_unflatten(cls, aux_data, children): + rc_scaled, zs_scaled = children + + if hasattr(rc_scaled, "shape") and hasattr(zs_scaled, "shape"): + 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 = cls._compute_scaling(xm, xn, scaling_type, scaling_factor) + + rc = rc_scaled / scaling + zs = zs_scaled / scaling + else: + rc = rc_scaled + zs = zs_scaled + + obj = object.__new__(cls) + obj._initialize_state( + rc, + zs, + aux_data["nfp"], + aux_data["mpol"], + aux_data["ntor"], + aux_data["ntheta"], + aux_data["nphi"], + aux_data["close"], + aux_data["range_torus"], + aux_data["scaling_type"], + aux_data["scaling_factor"], + ) + return obj + +tree_util.register_pytree_node(SurfaceRZFourier, + SurfaceRZFourier._tree_flatten, + SurfaceRZFourier._tree_unflatten) + #This class is based on simsopt classifier but translated to fit jax class SurfaceClassifier(): """ @@ -454,3 +890,8 @@ def signed_distance_from_surface_extras(xyz, surface): +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/coil_optimization/optimize_coils_and_nearaxis.py b/examples/coil_optimization/optimize_coils_and_nearaxis.py new file mode 100644 index 00000000..8be7ce54 --- /dev/null +++ b/examples/coil_optimization/optimize_coils_and_nearaxis.py @@ -0,0 +1,190 @@ +import os +number_of_processors_to_use = 1 # Parallelization, this should divide nfieldlines +os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' +from time import time +import jax.numpy as jnp +import matplotlib.pyplot as plt +from essos.coils import Coils, CreateEquallySpacedCurves +from essos.fields import BiotSavart +from pyqsc_jax.near_axis import near_axis +from essos.dynamics import Tracing +from essos.optimization import optimize_loss_function +from jax import vmap, jit +# In this exmple, `scipy.optimize.least_squares` is used, but any other optimizer, e.g. from +# `scipy.optimize.minimize` or `jaxopt`, can be used as well and may even be preferable. +from scipy.optimize import least_squares +from essos.losses import custom_loss + + +""" Creating starting coils and surface """ +N_COILS = 3; FOURIER_ORDER = 6; LARGE_R = 10; SMALL_R = 5.6; NFP = 3; N_SEGMENTS = 60; STELLSYM = True # Curve parameters +COIL_CURRENT = 1. # Amperes (optimization does not depend on current magnitude) +tolerance_optimization = 1e-8 +maximum_function_evaluations = 200 + + + +# Initialize Near-Axis field +rc=jnp.array([1.,0.01]) +zs=jnp.array([0.,0.01]) +etabar=-0.9 +field_nearaxis_initial = near_axis(rc=rc, zs=zs, etabar=etabar, nfp=NFP,order='r2') + +# Initialize coils +current_on_each_coil = 17.e5*field_nearaxis_initial.B0/NFP/2. +number_of_field_periods = NFP +major_radius_coils = field_nearaxis_initial.R0[0] +minor_radius_coils = major_radius_coils/2.0 +init_curves = CreateEquallySpacedCurves(n_curves=N_COILS, + order=FOURIER_ORDER, + R=major_radius_coils, r=minor_radius_coils, + n_segments=N_SEGMENTS, + nfp=number_of_field_periods, stellsym=STELLSYM) +init_coils = Coils(curves=init_curves, currents=jnp.array([current_on_each_coil]*N_COILS)) +init_field = BiotSavart(init_coils) + + +""" Setting the losses weights and targets """ +LENGTH_WEIGHT = 1.; LENGTH_TARGET = 4. +CURVATURE_WEIGHT = 1.; CURVATURE_TARGET = 6. +B_DIFFERENCE_WEIGHT = 1. +GRADB_DIFFERENCE_WEIGHT = 1. +IOTA_TARGET = 0.41 +IOTA_WEIGHT = 10. +R0_TARGET = field_nearaxis_initial.R0[0] +R0_WEIGHT = 10. + + + +""" Creating the loss functions """ +def near_axis_field_quantities(field_nearaxis): + Raxis = field_nearaxis.R0 + Zaxis = field_nearaxis.Z0 + phi = field_nearaxis.phi + Xaxis = Raxis*jnp.cos(phi) + Yaxis = Raxis*jnp.sin(phi) + points = jnp.array([Xaxis, Yaxis, Zaxis]) + B_nearaxis = field_nearaxis.B_axis.T + gradB_nearaxis = field_nearaxis.grad_B_axis.T + return points, B_nearaxis, gradB_nearaxis + + + +def loss_B_difference_coils_near_axis(field, field_nearaxis): + points, B_nearaxis, _ = near_axis_field_quantities(field_nearaxis) + B_coils = vmap(field.B)(points.T) + B_difference_loss = jnp.sum(jnp.abs(jnp.array(B_coils)-jnp.array(B_nearaxis))) + return B_difference_loss + +def loss_gradB_difference_coils_near_axis(field, field_nearaxis): + points, _, gradB_nearaxis = near_axis_field_quantities(field_nearaxis) + gradB_coils = vmap(field.dB_by_dX)(points.T) + gradB_difference_loss = jnp.sum(jnp.abs(jnp.array(gradB_coils)-jnp.array(gradB_nearaxis))) + return gradB_difference_loss + +def loss_iota_near_axis(field_nearaxis,iota_target=IOTA_TARGET): + return jnp.abs((field_nearaxis.iota - iota_target)) + +def loss_r0_near_axis(field_nearaxis, r0_target=R0_TARGET): + return jnp.abs((field_nearaxis.R0[0] - r0_target)) + +def loss_length(field,length_target=LENGTH_TARGET): + return jnp.mean(jnp.maximum(0, field.coils.length - length_target)) + +def loss_curvature(field,curvature_target=CURVATURE_TARGET): + return jnp.mean(jnp.maximum(0, field.coils.curvature - curvature_target)) + + +""" Defining custom losses """ +L_B_difference = custom_loss(loss_B_difference_coils_near_axis, "field", "field_nearaxis") +L_gradB_difference = custom_loss(loss_gradB_difference_coils_near_axis, "field", "field_nearaxis") +L_length = custom_loss(loss_length, "field") +L_curvature = custom_loss(loss_curvature, "field") +L_iota = custom_loss(loss_iota_near_axis, "field_nearaxis") +L_r0 = custom_loss(loss_r0_near_axis, "field_nearaxis") + + +""" Defining total loss + setting dependencies """ +L_total = B_DIFFERENCE_WEIGHT*L_B_difference + GRADB_DIFFERENCE_WEIGHT*L_gradB_difference + LENGTH_WEIGHT*L_length + CURVATURE_WEIGHT*L_curvature + IOTA_WEIGHT*L_iota+ R0_WEIGHT*L_r0 + + +L_total.dependencies = {"field": init_field, "field_nearaxis": field_nearaxis_initial} + +""" Optimizing the total loss """ +t_start = time() +res = least_squares(L_total, L_total.starting_dofs, L_total.grad, verbose=2, ftol=1e-5, gtol=1e-5, xtol=1e-14, max_nfev=maximum_function_evaluations) +t_end = time() + +print(f"\nOptimization took {t_end - t_start:.2f} seconds") +print("Initial loss:", L_total(L_total.starting_dofs)) +print("Loss after optimization:", L_total(res.x)) + +opt_field = L_total.dofs_to_pytree(res.x)["field"] +opt_coils = opt_field.coils + +opt_field_nearaxis = L_total.dofs_to_pytree(res.x)["field_nearaxis"] + + +B_difference_initial = loss_B_difference_coils_near_axis(init_field, field_nearaxis_initial) +gradB_difference_initial = loss_gradB_difference_coils_near_axis(init_field, field_nearaxis_initial) + +B_difference_optimized = loss_B_difference_coils_near_axis(opt_field, opt_field_nearaxis) +gradB_difference_optimized = loss_gradB_difference_coils_near_axis(opt_field, opt_field_nearaxis) + + +print(f'############################################') +print(f'Iota for initial near-axis: {field_nearaxis_initial.iota}') +print(f'Iota for optimized near-axis: {opt_field_nearaxis.iota}') +print(f'Maximum elongation for initial near-axis: {max(field_nearaxis_initial.elongation)}') +print(f'Maximum elongation for optimized near-axis: {max(opt_field_nearaxis.elongation)}') +print(f'Loss of B difference for initial near-axis: {B_difference_initial}') +print(f'Loss of B difference for optimized near-axis: {B_difference_optimized}') +print(f'Loss of gradB difference for initial near-axis: {gradB_difference_initial}') +print(f'Loss of gradB difference for optimized near-axis: {gradB_difference_optimized}') +print(f'Loss of R0 difference for initial near-axis: {loss_r0_near_axis(field_nearaxis_initial)}') +print(f'Loss of R0 difference for optimized near-axis: {loss_r0_near_axis(opt_field_nearaxis)}') + + +# Trace fieldlines +nfieldlines = 6 +num_steps = 1000 +tmax = 1.e-6 +trace_tolerance = 1e-7 + +R0_initial = jnp.linspace(field_nearaxis_initial.R0[0], 1.05*field_nearaxis_initial.R0[0], nfieldlines) +R0_optimized = jnp.linspace(opt_field_nearaxis.R0[0], 1.05*opt_field_nearaxis.R0[0], nfieldlines) +Z0 = jnp.zeros(nfieldlines) +phi0 = jnp.zeros(nfieldlines) +initial_xyz_initial = jnp.array([R0_initial*jnp.cos(phi0), R0_initial*jnp.sin(phi0), Z0]).T +initial_xyz_optimized = jnp.array([R0_optimized*jnp.cos(phi0), R0_optimized*jnp.sin(phi0), Z0]).T + +time0 = time() +tracing_initial = Tracing(field=init_field, model='FieldLineAdaptative', initial_conditions=initial_xyz_initial, + maxtime=tmax, times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance) +tracing_optimized = Tracing(field=opt_field, model='FieldLineAdaptative', initial_conditions=initial_xyz_optimized, + maxtime=tmax, times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance) +print(f"Tracing took {time()-time0:.2f} seconds") + +# Plot coils, before and after optimization +fig = plt.figure(figsize=(8, 4)) +ax1 = fig.add_subplot(121, projection='3d') +ax2 = fig.add_subplot(122, projection='3d') +init_coils.plot(ax=ax1, show=False) +field_nearaxis_initial.plot(ax=ax1, show=False, alpha=0.35) +tracing_initial.plot(ax=ax1, show=False) +opt_coils.plot(ax=ax2, show=False) +opt_field_nearaxis.plot(ax=ax2, show=False, alpha=0.35) +tracing_optimized.plot(ax=ax2, show=False) +plt.show() + +# # Save the coils to a json file +# coils_optimized.to_json("stellarator_coils.json") +# # Load the coils from a json file +# from essos.coils import Coils +# coils = Coils.from_json("stellarator_coils.json") + +# # Save results in vtk format to analyze in Paraview +# init_coils.to_vtk('coils_initial') +# opt_coils.to_vtk('coils_optimized') +# tracing_initial.to_vtk('trajectories_initial') +# tracing_optimized.to_vtk('trajectories_final') \ No newline at end of file diff --git a/examples/coil_optimization/optimize_coils_for_nearaxis.py b/examples/coil_optimization/optimize_coils_for_nearaxis.py new file mode 100644 index 00000000..f32eaa5f --- /dev/null +++ b/examples/coil_optimization/optimize_coils_for_nearaxis.py @@ -0,0 +1,168 @@ +import os +number_of_processors_to_use = 4 # Parallelization, this should divide nfieldlines +os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' +from time import time +import jax.numpy as jnp +import matplotlib.pyplot as plt +from essos.coils import Coils, CreateEquallySpacedCurves +from essos.fields import BiotSavart +from pyqsc_jax.near_axis import near_axis + +from essos.dynamics import Tracing +from essos.optimization import optimize_loss_function +from jax import vmap, jit +# In this exmple, `scipy.optimize.least_squares` is used, but any other optimizer, e.g. from +# `scipy.optimize.minimize` or `jaxopt`, can be used as well and may even be preferable. +from scipy.optimize import least_squares +from essos.losses import custom_loss + + +""" Creating starting coils and surface """ +N_COILS = 3; FOURIER_ORDER = 6; LARGE_R = 10; SMALL_R = 5.6; NFP = 3; N_SEGMENTS = 60; STELLSYM = True # Curve parameters +COIL_CURRENT = 1. # Amperes (optimization does not depend on current magnitude) +tolerance_optimization = 1e-8 +maximum_function_evaluations =200 + + + +# Initialize Near-Axis field +rc=jnp.array([1., 0.045]) +zs=jnp.array([0.,-0.045]) +etabar=-0.9 +field_nearaxis_initial = near_axis(rc=rc, zs=zs, etabar=etabar, nfp=NFP,order='r3') + +# Initialize coils +current_on_each_coil = 17.e5*field_nearaxis_initial.B0/NFP/2. +number_of_field_periods = NFP +major_radius_coils = field_nearaxis_initial.R0[0] +minor_radius_coils = major_radius_coils/2.0 +init_curves = CreateEquallySpacedCurves(n_curves=N_COILS, + order=FOURIER_ORDER, + R=major_radius_coils, r=minor_radius_coils, + n_segments=N_SEGMENTS, + nfp=number_of_field_periods, stellsym=STELLSYM) +init_coils = Coils(curves=init_curves, currents=jnp.array([current_on_each_coil]*N_COILS)) +init_field = BiotSavart(init_coils) + + +""" Setting the losses weights and targets """ +LENGTH_WEIGHT = 1.; LENGTH_TARGET = 4. +CURVATURE_WEIGHT = 1.; CURVATURE_TARGET = 6. +B_DIFFERENCE_WEIGHT = 1. +GRADB_DIFFERENCE_WEIGHT = 1. + + + +""" Creating the loss functions """ +def near_axis_field_quantities(field_nearaxis): + Raxis = field_nearaxis.R0 + Zaxis = field_nearaxis.Z0 + phi = field_nearaxis.phi + Xaxis = Raxis*jnp.cos(phi) + Yaxis = Raxis*jnp.sin(phi) + points = jnp.array([Xaxis, Yaxis, Zaxis]) + B_nearaxis = field_nearaxis.B_axis.T + gradB_nearaxis = field_nearaxis.grad_B_axis.T + return points, B_nearaxis, gradB_nearaxis + + +def loss_B_difference_coils_near_axis(field, field_nearaxis): + points, B_nearaxis, _ = near_axis_field_quantities(field_nearaxis) + B_coils = vmap(field.B)(points.T) + B_difference_loss = jnp.sum(jnp.abs(jnp.array(B_coils)-jnp.array(B_nearaxis))) + return B_difference_loss + +def loss_gradB_difference_coils_near_axis(field, field_nearaxis): + points, _, gradB_nearaxis = near_axis_field_quantities(field_nearaxis) + gradB_coils = vmap(field.dB_by_dX)(points.T) + gradB_difference_loss = jnp.sum(jnp.abs(jnp.array(gradB_coils)-jnp.array(gradB_nearaxis))) + return gradB_difference_loss + +def loss_length(field,length_target=LENGTH_TARGET): + return jnp.mean(jnp.maximum(0, field.coils.length - length_target)) + +def loss_curvature(field,curvature_target=CURVATURE_TARGET): + return jnp.mean(jnp.maximum(0, field.coils.curvature - curvature_target)) + + +""" Defining custom losses """ +L_B_difference = custom_loss(loss_B_difference_coils_near_axis, "field", field_nearaxis=field_nearaxis_initial) +L_gradB_difference = custom_loss(loss_gradB_difference_coils_near_axis, "field", field_nearaxis=field_nearaxis_initial) +L_length = custom_loss(loss_length, "field") +L_curvature = custom_loss(loss_curvature, "field") +""" Defining total loss + setting dependencies """ +L_total = B_DIFFERENCE_WEIGHT*L_B_difference + GRADB_DIFFERENCE_WEIGHT*L_gradB_difference + LENGTH_WEIGHT*L_length + CURVATURE_WEIGHT*L_curvature + + +L_total.dependencies = {"field": init_field, "field_nearaxis": field_nearaxis_initial} + +""" Optimizing the total loss """ +t_start = time() +res = least_squares(L_total, L_total.starting_dofs, L_total.grad, verbose=2, ftol=1e-5, gtol=1e-5, xtol=1e-14, max_nfev=200) +t_end = time() + +print(f"\nOptimization took {t_end - t_start:.2f} seconds") +print("Initial loss:", L_total(L_total.starting_dofs)) +print("Loss after optimization:", L_total(res.x)) + +opt_field = L_total.dofs_to_pytree(res.x)["field"] +opt_coils = opt_field.coils + + +B_difference_initial = loss_B_difference_coils_near_axis(init_field, field_nearaxis_initial) +gradB_difference_initial = loss_gradB_difference_coils_near_axis(init_field, field_nearaxis_initial) + +B_difference_optimized = loss_B_difference_coils_near_axis(opt_field, field_nearaxis_initial) +gradB_difference_optimized = loss_gradB_difference_coils_near_axis(opt_field, field_nearaxis_initial) + + +print(f'############################################') +print(f'Loss of B difference for initial near-axis: {B_difference_initial}') +print(f'Loss of B difference for optimized near-axis: {B_difference_optimized}') +print(f'Loss of gradB difference for initial near-axis: {gradB_difference_initial}') +print(f'Loss of gradB difference for optimized near-axis: {gradB_difference_optimized}') + + +# Trace fieldlines +nfieldlines = 6 +num_steps = 1000 +tmax = 1.e-6 +trace_tolerance = 1e-7 + +R0_initial = jnp.linspace(field_nearaxis_initial.R0[0], 1.05*field_nearaxis_initial.R0[0], nfieldlines) +R0_optimized = jnp.linspace(field_nearaxis_initial.R0[0], 1.05*field_nearaxis_initial.R0[0], nfieldlines) +Z0 = jnp.zeros(nfieldlines) +phi0 = jnp.zeros(nfieldlines) +initial_xyz_initial = jnp.array([R0_initial*jnp.cos(phi0), R0_initial*jnp.sin(phi0), Z0]).T +initial_xyz_optimized = jnp.array([R0_optimized*jnp.cos(phi0), R0_optimized*jnp.sin(phi0), Z0]).T + +time0 = time() +tracing_initial = Tracing(field=init_field, model='FieldLineAdaptative', initial_conditions=initial_xyz_initial, + maxtime=tmax, times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance) +tracing_optimized = Tracing(field=opt_field, model='FieldLineAdaptative', initial_conditions=initial_xyz_optimized, + maxtime=tmax, times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance) +print(f"Tracing took {time()-time0:.2f} seconds") + +# Plot coils, before and after optimization +fig = plt.figure(figsize=(8, 4)) +ax1 = fig.add_subplot(121, projection='3d') +ax2 = fig.add_subplot(122, projection='3d') +init_coils.plot(ax=ax1, show=False) +field_nearaxis_initial.plot(ax=ax1, show=False, alpha=0.35) +tracing_initial.plot(ax=ax1, show=False) +opt_coils.plot(ax=ax2, show=False) +field_nearaxis_initial.plot(ax=ax2, show=False, alpha=0.35) +tracing_optimized.plot(ax=ax2, show=False) +plt.show() + +# # Save the coils to a json file +# coils_optimized.to_json("stellarator_coils.json") +# # Load the coils from a json file +# from essos.coils import Coils +# coils = Coils.from_json("stellarator_coils.json") + +# # Save results in vtk format to analyze in Paraview +# init_coils.to_vtk('coils_initial') +# opt_coils.to_vtk('coils_optimized') +# tracing_initial.to_vtk('trajectories_initial') +# tracing_optimized.to_vtk('trajectories_final') \ No newline at end of file diff --git a/examples/coil_optimization/optimize_coils_particle_confinement_fullorbit.py b/examples/coil_optimization/optimize_coils_particle_confinement_fullorbit.py new file mode 100644 index 00000000..d68d3a13 --- /dev/null +++ b/examples/coil_optimization/optimize_coils_particle_confinement_fullorbit.py @@ -0,0 +1,156 @@ + +import os +number_of_processors_to_use = 1 # Parallelization, this should divide nparticles +os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' +from time import time +import jax.numpy as jnp +import matplotlib.pyplot as plt +from essos.dynamics import Particles, Tracing +from essos.coils import Coils, CreateEquallySpacedCurves +from jax import vmap, jit +# In this exmple, `scipy.optimize.least_squares` is used, but any other optimizer, e.g. from +# `scipy.optimize.minimize` or `jaxopt`, can be used as well and may even be preferable. +from scipy.optimize import least_squares +from essos.losses import custom_loss +from essos.fields import BiotSavart + + +# Particle optimization parameters +# Optimization parameters +NPARTICLES = number_of_processors_to_use*10 +MAXTIME_TRACING = 1e-4 +NUMBER_COILS_PER_HALF_FIELD_PERIOD = 3 +NUMBER_OF_FIELD_PERIODS = 2 +MODEL = 'FullOrbit_Boris' +TIMESTEP=1.e-14 +TRACE_TOLERANCE=1e-8 +NUM_STEPS=1000 + + +NPARTICLES_PLOT = number_of_processors_to_use*10 +MAXTIME_TRACING_PLOT = 1e-4 + +""" Creating starting coils and surface """ +N_COILS = 3 +FOURIER_ORDER = 6 +LARGE_R = 7.74 +SMALL_R = 4.5 +NFP = 2 +N_SEGMENTS = 60 +STELLSYM = True # Curve parameters +COIL_CURRENT = 1.84e7 # Amperes (optimization does not depend on current magnitude) +MAXIMUM_FUNCTION_EVALUATIONS =200 + +# Initialize coils +init_curves = CreateEquallySpacedCurves(n_curves=N_COILS, + order=FOURIER_ORDER, + R=LARGE_R, r=SMALL_R, + n_segments=N_SEGMENTS, + nfp=NFP, stellsym=STELLSYM) +init_coils = Coils(curves=init_curves, currents=jnp.array([COIL_CURRENT]*N_COILS)) +init_field = BiotSavart(init_coils) + +""" Setting the losses weights and targets """ +LENGTH_WEIGHT = 1.; LENGTH_TARGET = 31. +CURVATURE_WEIGHT = 1.; CURVATURE_TARGET = 0.4 +BAXIS_WEIGHT = 1.; BAXIS_TARGET = 5.7 +RADIAL_DRIFT_WEIGHT = 1. + +def loss_particle_radial_drift(field, particles, timestep=1.e-8, maxtime=1e-5, num_steps=300, trace_tolerance=1e-5, model='GuidingCenterAdaptative',boundary=None): + particles.to_full_orbit(field) + tracing = Tracing(field=field, model=model, particles=particles, maxtime=maxtime, + timestep=timestep,times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance,boundary=boundary) + xyz = tracing.trajectories[:,:, :3] + R_axis=field.r_axis + Z_axis=field.z_axis + #Ideally here one would differentiate in time through diffrax !TODO + r_cross=jnp.sqrt(jnp.square(jnp.sqrt(jnp.square(xyz[:,0])+jnp.square(xyz[:,1]))-R_axis+1.e-12)+jnp.square(xyz[:,2]-Z_axis+1.e-12)) + v_r_cross=jnp.diff(r_cross,axis=1)#/tracing.times_to_trace*tracing.maxtime + return (jnp.sum(jnp.square(jnp.average(v_r_cross,axis=1)))) + +def normB_axis(field, npoints=15): + R_axis=field.r_axis + phi_array = jnp.linspace(0, 2 * jnp.pi, npoints) + B_axis = vmap(lambda phi: field.AbsB(jnp.array([R_axis * jnp.cos(phi), R_axis * jnp.sin(phi), 0])))(phi_array) + return B_axis + +def loss_normB_axis_average(field,npoints=15, target_B=BAXIS_TARGET): + B_axis = normB_axis(field, npoints) + return jnp.abs(jnp.average(B_axis)-target_B) + +def loss_length(field,length_target=LENGTH_TARGET): + return jnp.mean(jnp.maximum(0, field.coils.length - length_target)) + +def loss_curvature(field,curvature_target=CURVATURE_TARGET): + return jnp.mean(jnp.maximum(0, field.coils.curvature - curvature_target)) + +# Initialize particles +phi_array = jnp.linspace(0, 2*jnp.pi, NPARTICLES) +initial_xyz=jnp.array([LARGE_R*jnp.cos(phi_array), LARGE_R*jnp.sin(phi_array), 0*phi_array]).T +particles = Particles(initial_xyz=initial_xyz) + + +""" Defining custom losses """ +L_radial_drift= custom_loss(loss_particle_radial_drift, "field", particles=particles, timestep=TIMESTEP, maxtime=MAXTIME_TRACING, num_steps=NUM_STEPS, trace_tolerance=TRACE_TOLERANCE, model=MODEL) +L_B_axis= custom_loss(loss_normB_axis_average, "field") +L_length = custom_loss(loss_length, "field") +L_curvature = custom_loss(loss_curvature, "field") +""" Defining total loss + setting dependencies """ +L_total = RADIAL_DRIFT_WEIGHT*L_radial_drift + L_B_axis + LENGTH_WEIGHT*L_length + CURVATURE_WEIGHT*L_curvature + + +L_total.dependencies = {"field": init_field} + +""" Optimizing the total loss """ +t_start = time() +res = least_squares(L_total, L_total.starting_dofs, L_total.grad, verbose=2, ftol=1e-5, gtol=1e-5, xtol=1e-14, max_nfev=MAXIMUM_FUNCTION_EVALUATIONS) +t_end = time() + +print(f"\nOptimization took {t_end - t_start:.2f} seconds") +print("Initial loss:", L_total(L_total.starting_dofs)) +print("Loss after optimization:", L_total(res.x)) + +opt_field = L_total.dofs_to_pytree(res.x)["field"] +opt_coils = opt_field.coils + +""" Plotting results """ +phi_array_plot = jnp.linspace(0, 2*jnp.pi, NPARTICLES_PLOT) +initial_xyz_plot=jnp.array([LARGE_R*jnp.cos(phi_array_plot), LARGE_R*jnp.sin(phi_array_plot), 0*phi_array_plot]).T +particles_plot = Particles(initial_xyz=initial_xyz_plot) +particles_plot.to_full_orbit(opt_field) +tracing_initial = Tracing(field=init_field, particles=particles_plot, maxtime=MAXTIME_TRACING_PLOT, model=MODEL, times_to_trace=NUM_STEPS, timestep=TIMESTEP, atol=TRACE_TOLERANCE, rtol=TRACE_TOLERANCE) +tracing_optimized = Tracing(field=opt_field, particles=particles_plot, maxtime=MAXTIME_TRACING_PLOT, model=MODEL, times_to_trace=NUM_STEPS, timestep=TIMESTEP, atol=TRACE_TOLERANCE, rtol=TRACE_TOLERANCE) + +# Plot trajectories, before and after optimization +fig = plt.figure(figsize=(9, 8)) +ax1 = fig.add_subplot(221, projection='3d') +ax2 = fig.add_subplot(222, projection='3d') +ax3 = fig.add_subplot(223) +ax4 = fig.add_subplot(224) + +init_coils.plot(ax=ax1, show=False) +tracing_initial.plot(ax=ax1, show=False) +for i, trajectory in enumerate(tracing_initial.trajectories): + ax3.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') + +ax3.set_xlabel('R (m)');ax3.set_ylabel('Z (m)');#ax3.legend() +opt_coils.plot(ax=ax2, show=False) +tracing_optimized.plot(ax=ax2, show=False) +for i, trajectory in enumerate(tracing_optimized.trajectories): + ax4.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') + +ax4.set_xlabel('R (m)');ax4.set_ylabel('Z (m)');#ax4.legend() +plt.tight_layout() +plt.show() + +# # Save the coils to a json file +# coils_optimized.to_json("stellarator_coils.json") +# # Load the coils from a json file +# from essos.coils import Coils +# coils = Coils.from_json("stellarator_coils.json") + +# # Save results in vtk format to analyze in Paraview +# tracing_initial.to_vtk('trajectories_initial') +# tracing_optimized.to_vtk('trajectories_final') +# coils_initial.to_vtk('coils_initial') +# coils_optimized.to_vtk('coils_optimized') \ No newline at end of file diff --git a/examples/coil_optimization/optimize_coils_particle_confinement_guidingcenter.py b/examples/coil_optimization/optimize_coils_particle_confinement_guidingcenter.py new file mode 100644 index 00000000..0c55d5eb --- /dev/null +++ b/examples/coil_optimization/optimize_coils_particle_confinement_guidingcenter.py @@ -0,0 +1,159 @@ + +import os +number_of_processors_to_use = 1 # Parallelization, this should divide nparticles +os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' +from time import time +import jax.numpy as jnp +import matplotlib.pyplot as plt +from essos.dynamics import Particles, Tracing +from essos.coils import Coils, CreateEquallySpacedCurves +from jax import vmap, jit +# In this exmple, `scipy.optimize.least_squares` is used, but any other optimizer, e.g. from +# `scipy.optimize.minimize` or `jaxopt`, can be used as well and may even be preferable. +from scipy.optimize import least_squares +from essos.losses import custom_loss +from essos.fields import BiotSavart + +# Particle optimization parameters +# Optimization parameters +NPARTICLES = number_of_processors_to_use*10 +MAXTIME_TRACING = 1e-4 +NUMBER_COILS_PER_HALF_FIELD_PERIOD = 3 +NUMBER_OF_FIELD_PERIODS = 2 +MODEL = 'GuidingCenterAdaptative' +TIMESTEP=1.e-14 +TRACE_TOLERANCE=1e-8 +NUM_STEPS=1000 + +NPARTICLES_PLOT = number_of_processors_to_use*10 +MAXTIME_TRACING_PLOT = 1e-4 + + + +""" Creating starting coils and surface """ +N_COILS = 3 +FOURIER_ORDER = 6 +LARGE_R = 7.74 +SMALL_R = 4.5 +NFP = 2 +N_SEGMENTS = 60 +STELLSYM = True # Curve parameters +COIL_CURRENT = 1.84e7 # Amperes (optimization does not depend on current magnitude) +MAXIMUM_FUNCTION_EVALUATIONS =200 + +# Initialize coils +init_curves = CreateEquallySpacedCurves(n_curves=N_COILS, + order=FOURIER_ORDER, + R=LARGE_R, r=SMALL_R, + n_segments=N_SEGMENTS, + nfp=NFP, stellsym=STELLSYM, + scaling_type="L2", + scaling_factor=1.2, + scale_fixed=1.0) +init_coils = Coils(curves=init_curves, currents=jnp.array([COIL_CURRENT]*N_COILS)) +init_field = BiotSavart(init_coils) + +""" Setting the losses weights and targets """ +LENGTH_WEIGHT = 1.; LENGTH_TARGET = 31. +CURVATURE_WEIGHT = 1.; CURVATURE_TARGET = 0.4 +BAXIS_WEIGHT = 1.; BAXIS_TARGET = 5.7 +RADIAL_DRIFT_WEIGHT = 1. + +def loss_particle_radial_drift(field, particles, timestep=1.e-8, maxtime=1e-4, num_steps=300, trace_tolerance=1e-5, model='GuidingCenterAdaptative',boundary=None): + particles.to_full_orbit(field) + tracing = Tracing(field=field, model=model, particles=particles, maxtime=maxtime, + timestep=timestep,times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance,boundary=boundary) + xyz = tracing.trajectories[:,:, :3] + R_axis=field.r_axis + Z_axis=field.z_axis + #Ideally here one would differentiate in time through diffrax !TODO + r_cross=jnp.sqrt(jnp.square(jnp.sqrt(jnp.square(xyz[:,:,0])+jnp.square(xyz[:,:,1]))-R_axis+1.e-12)+jnp.square(xyz[:,:,2]-Z_axis+1.e-12)) + return jnp.average(r_cross[:,-1]) # Average radial end position of all particles at the end of the tracing, ideally this should be zero for perfect confinement + +def normB_axis(field, npoints=15): + R_axis=field.r_axis + phi_array = jnp.linspace(0, 2 * jnp.pi, npoints) + B_axis = vmap(lambda phi: field.AbsB(jnp.array([R_axis * jnp.cos(phi), R_axis * jnp.sin(phi), 0])))(phi_array) + return B_axis + +def loss_normB_axis_average(field,npoints=15, target_B=BAXIS_TARGET): + B_axis = normB_axis(field, npoints) + return jnp.abs(jnp.average(B_axis)-target_B) + +def loss_length(field,length_target=LENGTH_TARGET): + return jnp.mean(jnp.maximum(0, field.coils.length - length_target)) + +def loss_curvature(field,curvature_target=CURVATURE_TARGET): + return jnp.mean(jnp.maximum(0, field.coils.curvature - curvature_target)) + + + +# Initialize particles +phi_array = jnp.linspace(0, 2*jnp.pi, NPARTICLES) +initial_xyz=jnp.array([LARGE_R*jnp.cos(phi_array), LARGE_R*jnp.sin(phi_array), 0*phi_array]).T +particles = Particles(initial_xyz=initial_xyz) + + +""" Defining custom losses """ +L_radial_drift= custom_loss(loss_particle_radial_drift, "field", particles=particles, timestep=TIMESTEP, maxtime=MAXTIME_TRACING, num_steps=NUM_STEPS, trace_tolerance=TRACE_TOLERANCE, model=MODEL) +L_B_axis= custom_loss(loss_normB_axis_average, "field") +L_length = custom_loss(loss_length, "field") +L_curvature = custom_loss(loss_curvature, "field") +""" Defining total loss + setting dependencies """ +L_total = RADIAL_DRIFT_WEIGHT*L_radial_drift + L_B_axis + LENGTH_WEIGHT*L_length + CURVATURE_WEIGHT*L_curvature + +L_total.dependencies = {"field": init_field} + +""" Optimizing the total loss """ +t_start = time() +res = least_squares(L_total, L_total.starting_dofs, L_total.grad, verbose=2, ftol=1e-5, gtol=1e-5, xtol=1e-14, max_nfev=MAXIMUM_FUNCTION_EVALUATIONS) +t_end = time() + +print(f"\nOptimization took {t_end - t_start:.2f} seconds") +print("Initial loss:", L_total(L_total.starting_dofs)) +print("Loss after optimization:", L_total(res.x)) + +opt_field = L_total.dofs_to_pytree(res.x)["field"] +opt_coils = opt_field.coils + +""" Plotting results """ +phi_array_plot = jnp.linspace(0, 2*jnp.pi, NPARTICLES_PLOT) +initial_xyz_plot=jnp.array([LARGE_R*jnp.cos(phi_array_plot), LARGE_R*jnp.sin(phi_array_plot), 0*phi_array_plot]).T +particles_plot = Particles(initial_xyz=initial_xyz_plot) +particles_plot.to_full_orbit(opt_field) +tracing_initial = Tracing(field=init_field, particles=particles_plot, maxtime=MAXTIME_TRACING_PLOT, model=MODEL, times_to_trace=NUM_STEPS, timestep=TIMESTEP, atol=TRACE_TOLERANCE, rtol=TRACE_TOLERANCE) +tracing_optimized = Tracing(field=opt_field, particles=particles_plot, maxtime=MAXTIME_TRACING_PLOT, model=MODEL, times_to_trace=NUM_STEPS, timestep=TIMESTEP, atol=TRACE_TOLERANCE, rtol=TRACE_TOLERANCE) + +# Plot trajectories, before and after optimization +fig = plt.figure(figsize=(9, 8)) +ax1 = fig.add_subplot(221, projection='3d') +ax2 = fig.add_subplot(222, projection='3d') +ax3 = fig.add_subplot(223) +ax4 = fig.add_subplot(224) + +init_coils.plot(ax=ax1, show=False) +tracing_initial.plot(ax=ax1, show=False) +for i, trajectory in enumerate(tracing_initial.trajectories): + ax3.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') + +ax3.set_xlabel('R (m)');ax3.set_ylabel('Z (m)');#ax3.legend() +opt_coils.plot(ax=ax2, show=False) +tracing_optimized.plot(ax=ax2, show=False) +for i, trajectory in enumerate(tracing_optimized.trajectories): + ax4.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') + +ax4.set_xlabel('R (m)');ax4.set_ylabel('Z (m)');#ax4.legend() +plt.tight_layout() +plt.show() + +# # Save the coils to a json file +# coils_optimized.to_json("stellarator_coils.json") +# # Load the coils from a json file +# from essos.coils import Coils_from_json +# coils = Coils_from_json("stellarator_coils.json") + +# # Save results in vtk format to analyze in Paraview +# tracing_initial.to_vtk('trajectories_initial') +# tracing_optimized.to_vtk('trajectories_final') +# coils_initial.to_vtk('coils_initial') +# coils_optimized.to_vtk('coils_optimized') diff --git a/examples/coil_optimization/optimize_coils_vmec_surface.py b/examples/coil_optimization/optimize_coils_vmec_surface.py new file mode 100644 index 00000000..65a38ea9 --- /dev/null +++ b/examples/coil_optimization/optimize_coils_vmec_surface.py @@ -0,0 +1,86 @@ +import os +from time import time +import jax.numpy as jnp +import matplotlib.pyplot as plt + +from essos.coils import Coils, CreateEquallySpacedCurves +from essos.fields import BiotSavart +from essos.surfaces import SurfaceRZFourier, BdotN_over_B +from essos.losses import custom_loss + +# In this exmple, `scipy.optimize.least_squares` is used, but any other optimizer, e.g. from +# `scipy.optimize.minimize` or `jaxopt`, can be used as well and may even be preferable. +from scipy.optimize import least_squares + +input_filepath = os.path.join(os.path.dirname(__file__), "..", "input_files") +vmec_input = os.path.join(input_filepath, 'wout_LandremanPaul2021_QA_reactorScale_lowres.nc') + +""" Creating starting coils and surface """ +N_COILS = 3; FOURIER_ORDER = 3; LARGE_R = 10; SMALL_R = 5.6; NFP = 2; N_SEGMENTS = 45; STELLSYM = True # Curve parameters +COIL_CURRENT = 1. # Amperes (optimization does not depend on current magnitude) + +init_curves = CreateEquallySpacedCurves(N_COILS, FOURIER_ORDER, LARGE_R, SMALL_R, n_segments=N_SEGMENTS, nfp=NFP, stellsym=STELLSYM) +init_coils = Coils(curves=init_curves, currents=[COIL_CURRENT]*N_COILS) +init_field = BiotSavart(init_coils) +surface = SurfaceRZFourier.from_wout_file(vmec_input, s=1, ntheta=30, nphi=30, range_torus='half period') + +""" Setting the losses weights and targets """ +LENGTH_WEIGHT = 1.; LENGTH_TARGET = 32. +CURVATURE_WEIGHT = 1.; CURVATURE_TARGET = 0.1 +NORMAL_FIELD_WEIGHT = 1. + +""" Creating the loss functions """ +def loss(field, surface): + return jnp.sum(jnp.abs(BdotN_over_B(surface, field))) + +def loss_length(field): + return jnp.mean(jnp.maximum(0, field.coils.length - LENGTH_TARGET)) + +def loss_curvature(field): + return jnp.mean(jnp.maximum(0, field.coils.curvature - CURVATURE_TARGET)) + +""" Defining custom losses """ +L_normal_field = custom_loss(loss, "field", surface=surface) +L_length = custom_loss(loss_length, "field") +L_curvature = custom_loss(loss_curvature, "field") + +""" Defining total loss + setting dependencies """ +L_total = NORMAL_FIELD_WEIGHT*L_normal_field + LENGTH_WEIGHT*L_length + CURVATURE_WEIGHT*L_curvature +L_total.dependencies = {"field": init_field} + +""" Optimizing the total loss """ +t_start = time() +res = least_squares(L_total, L_total.starting_dofs, L_total.grad, verbose=2, ftol=1e-5, gtol=1e-5, xtol=1e-14, max_nfev=200) +t_end = time() + +print(f"\nOptimization took {t_end - t_start:.2f} seconds") +print("Initial loss:", L_total(L_total.starting_dofs)) +print("Loss after optimization:", L_total(res.x)) + +opt_field = L_total.dofs_to_pytree(res.x)["field"] +opt_coils = opt_field.coils + +fig = plt.figure(figsize=(8, 4)) + +ax1 = fig.add_subplot(121, projection='3d') +init_coils.plot(ax=ax1, show=False) +surface.plot(ax=ax1, show=False) +ax2 = fig.add_subplot(122, projection='3d') +opt_coils.plot(ax=ax2, show=False) +surface.plot(ax=ax2, show=False) +plt.tight_layout() +plt.show() + +EXPORT = False +if EXPORT: + output_filepath = os.path.join(os.path.dirname(__file__), "output") + + """ Save the coils to a json file """ + init_coils.to_json(os.path.join(output_filepath, "init_coils_vmec_surface.json")) + opt_coils.to_json(os.path.join(output_filepath, "opt_coils_vmec_surface.json")) + + """ Save results in vtk format to analyze in Paraview """ + surface.to_vtk(os.path.join(output_filepath, "init_surface_vmec_surface.json"), field=init_field) + surface.to_vtk(os.path.join(output_filepath, "final_surface_vmec_surface.json"), field=opt_field) + init_coils.to_vtk(os.path.join(output_filepath, "init_coils_vmec_surface.json")) + opt_coils.to_vtk(os.path.join(output_filepath, "opt_coils_vmec_surface.json")) \ No newline at end of file diff --git a/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py b/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py new file mode 100644 index 00000000..72818bbf --- /dev/null +++ b/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py @@ -0,0 +1,226 @@ +import os +from time import time +import jax +import jax.numpy as jnp +import matplotlib.pyplot as plt +from essos.optimization import optimize_loss_function +from essos.coils import Coils, CreateEquallySpacedCurves +from essos.fields import BiotSavart +from essos.surfaces import SurfaceRZFourier, BdotN_over_B +from essos.losses import custom_loss, base_loss + +import essos.augmented_lagrangian as alm +from functools import partial +# In this exmple, `scipy.optimize.least_squares` is used for the normal optimization, but any other optimizer, e.g. from +# `scipy.optimize.minimize` or `jaxopt`, can be used as well and may even be preferable. +from scipy.optimize import least_squares + +# Optimization parameters +maximum_function_evaluations=100 + +input_filepath = os.path.join(os.path.dirname(__file__), "..", "input_files") +vmec_input = os.path.join(input_filepath, 'wout_LandremanPaul2021_QA_reactorScale_lowres.nc') +surface = SurfaceRZFourier.from_wout_file(vmec_input, s=1, ntheta=32, nphi=32, range_torus='half period') + + +""" Creating starting coils and surface """ +N_COILS = 4 +FOURIER_ORDER = 3 +LARGE_R = 10. +SMALL_R = 5.7 +NFP = 2 +N_SEGMENTS = 30 +STELLSYM = True # Curve parameters +COIL_CURRENT = 1. # Amperes (optimization does not depend on current magnitude) + + +""" Setting the losses weights and targets """ +LENGTH_WEIGHT = 1. +LENGTH_TARGET = 40. +CURVATURE_WEIGHT = 1. +CURVATURE_TARGET = 0.5 +NORMAL_FIELD_WEIGHT = 1. +BdotN_Target_tol=1.e-6 + +EXPORT = False + +init_curves = CreateEquallySpacedCurves(N_COILS, FOURIER_ORDER, LARGE_R, SMALL_R, n_segments=N_SEGMENTS, nfp=NFP, stellsym=STELLSYM) +init_coils = Coils(curves=init_curves, currents=[COIL_CURRENT]*N_COILS) +init_field = BiotSavart(init_coils) +init_surface=surface + + +""" Creating the loss functions """ +def loss(field, surface): + return jnp.sum(jnp.abs(BdotN_over_B(surface, field))) + +def BdotN_constraint(field,surface,target_tol=1.e-6): + bdotn_over_b = BdotN_over_B(surface, field) + bdotn_over_b_loss = jnp.sqrt(jnp.sum(jnp.maximum(jnp.square(bdotn_over_b)-target_tol,0.0))) + return bdotn_over_b_loss + +def loss_length_constraint(field): + return jnp.maximum(0, field.coils.length - LENGTH_TARGET) + +def loss_curvature_contraint(field): + return jnp.maximum(0, field.coils.curvature - CURVATURE_TARGET) + #return jnp.mean(jnp.maximum(0, (field.coils.curvature - CURVATURE_TARGET)/CURVATURE_TARGET)) + +def loss_length(field): + return jnp.mean(jnp.maximum(0, field.coils.length - LENGTH_TARGET)) + +def loss_curvature(field): + return jnp.mean(jnp.maximum(0, field.coils.curvature - CURVATURE_TARGET)) + + +""" Defining custom losses """ +L_normal_field = custom_loss(loss, "field", surface=surface) +L_normal_field_constraint = custom_loss(BdotN_constraint, "field", surface=surface) +L_length_constraint = custom_loss(loss_length_constraint, "field") +L_curvature_constraint = custom_loss(loss_curvature_contraint, "field") +L_length = custom_loss(loss_length, "field") +L_curvature = custom_loss(loss_curvature, "field") + +""" Defining total loss + setting dependencies """ +L_normal_field.dependencies = {"field": init_field} +L_length_constraint.dependencies = {"field": init_field} +L_curvature_constraint.dependencies = {"field": init_field} +L_normal_field_constraint.dependencies = {"field": init_field} +L_length.dependencies = {"field": init_field} +L_curvature.dependencies = {"field": init_field} + + +# Create the constraints +penalty = 1.0 #Intial penalty values +multiplier=0. #Initial lagrange multiplier values +#Initializing first tolerances for the inner minimisation loop iteration +omega=1./penalty +eta=1./penalty**0.1 +sq_grad=0.0 #Initial square gradient parameter value for Mu adaptative +model_lagrangian='Standard' #Use standard augmented lagragian suitable for bounded optimizers +#Since we are using LBFGS-B from jaxopt, model_mu will be updated with tolerances so we do not need to difinte the model +model_mu='Tolerance' + + +beta=2. #penalty update parameter +mu_max=1.e4 #Maximum penalty parameter allowed +alpha=0.99 #These are parameters only used if gradient descent and adaaptative mu +gamma=1.e-2 +epsilon=1.e-8 +omega_tol=1.e-7 #desired grad_tolerance, associated with grad of lagrangian to main parameters +eta_tol=1.e-7 #desired contraint tolerance, associated with variation of contraints + +#curvature_constraint=alm.ScaledConstraint(alm.eq(loss_curvature_contraint,model_lagrangian=model_lagrangian, multiplier=multiplier,penalty=penalty,omega=omega,sq_grad=sq_grad)) +#length_constraint=alm.ScaledConstraint(alm.eq(loss_length_constraint,model_lagrangian=model_lagrangian, multiplier=multiplier,penalty=penalty,omega=omega,sq_grad=sq_grad)) +#field_constraint=alm.ScaledConstraint(alm.eq(BdotN_constraint,model_lagrangian=model_lagrangian, multiplier=multiplier,penalty=penalty,omega=omega,sq_grad=sq_grad)) +curvature_constraint=alm.eq(loss_curvature_contraint,model_lagrangian=model_lagrangian, multiplier=multiplier,penalty=penalty,omega=omega,sq_grad=sq_grad) +length_constraint=alm.eq(loss_length_constraint,model_lagrangian=model_lagrangian, multiplier=multiplier,penalty=penalty,omega=omega,sq_grad=sq_grad) +field_constraint=alm.eq(BdotN_constraint,model_lagrangian=model_lagrangian, multiplier=multiplier,penalty=penalty,omega=omega,sq_grad=sq_grad) + + +C_normal_field_constraint = alm.SelectiveConstraint(field_constraint, "field", surface=surface, target_tol=BdotN_Target_tol) +C_length_constraint = alm.SelectiveConstraint(length_constraint, "field") +C_curvature_constraint = alm.SelectiveConstraint(curvature_constraint, "field") + + + +C_Total_constraint = alm.combine(C_normal_field_constraint, C_length_constraint, C_curvature_constraint) + + +C_normal_field_constraint.dependencies = {"field": init_field} +C_length_constraint.dependencies = {"field": init_field} +C_curvature_constraint.dependencies = {"field": init_field} +C_Total_constraint.dependencies = {"field": init_field} + + +#If loss=cost_function(x) is not prescribed, f(x)=0 is considered, uncomment second line to use B dot N as a loss and not a constraint +ALM=alm.ALM_model_jaxopt_lbfgsb(constraints=C_Total_constraint,model_lagrangian=model_lagrangian,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) + +#Initializing lagrange multipliers +lagrange_params=C_Total_constraint.init(init_field.dofs) +#parameters are a tuple of the primal/main optimisation parameters and the lagrange multipliers +params = init_field.dofs, lagrange_params +#This is just to initialize an empty state for the lagrange multiplier update and get some information +lag_state,grad,info=ALM.init(params) + + +""" Optimizing with alm""" +t_start = time() + +i=0 +while i<=maximum_function_evaluations and (jnp.linalg.norm(grad[0])>omega_tol or alm.norm_constraints(info[2])>eta_tol): + #One step of ALM optimization + params, lag_state,grad,info = ALM.update(params,lag_state,grad,info) + #if i % 5 == 0: + #print(f'i: {i}, loss f: {info[0]:g}, infeasibility: {alm.total_infeasibility(info[1]):g}') + print(f'i: {i}, loss f: {info[0]:g},loss L: {info[1]:g}, infeasibility: {alm.total_infeasibility(info[2]):g}') + #print('lagrange',params[1]) + i=i+1 + +t_end = time() + + +opt_field_alm = C_normal_field_constraint.dofs_to_pytree(params[0])[0] +opt_coils_alm = opt_field_alm.coils + + +""" Defining total loss for nornmal optimization""" +L_total = NORMAL_FIELD_WEIGHT*L_normal_field+ LENGTH_WEIGHT*L_length + CURVATURE_WEIGHT*L_curvature +L_total.dependencies = {"field": init_field} + +""" Optimizing the total loss """ +t_start = time() +res = least_squares(L_total, L_total.starting_dofs, L_total.grad, verbose=2, ftol=1e-5, gtol=1e-5, xtol=1e-14, max_nfev=maximum_function_evaluations) +t_end = time() + +print(f"\nOptimization took {t_end - t_start:.2f} seconds") +print("Initial loss:", L_total(L_total.starting_dofs)) +print("Loss after optimization:", L_total(res.x)) + +opt_field = L_total.dofs_to_pytree(res.x)["field"] +opt_coils = opt_field.coils + + +print(f"\nOptimization took {t_end - t_start:.2f} seconds") +print("Initial B dot N:", jnp.max(BdotN_over_B(surface, init_field))) +print("B dot N after optimization:", jnp.max(BdotN_over_B(surface, opt_field))) +print("B dot N after optimization alm:", jnp.max(BdotN_over_B(surface, opt_field_alm))) +print("Initial curvature :", jnp.average(init_field.coils.curvature,axis=0)) +print("Curvature after optimization:",jnp.average(opt_field.coils.curvature,axis=0)) +print("Curvature after optimization alm:",jnp.average(opt_field_alm.coils.curvature,axis=0)) +print("Curvature target:",CURVATURE_TARGET) +print("Initial length :", init_field.coils.length) +print("Length after optimization:",opt_field.coils.length) +print("Length after optimization alm:",opt_field_alm.coils.length) +print("Length target:",LENGTH_TARGET) + + + +fig = plt.figure(figsize=(8, 4)) + +ax1 = fig.add_subplot(131, projection='3d') +init_coils.plot(ax=ax1, show=False,label='Initial coils') +surface.plot(ax=ax1, show=False) +ax2 = fig.add_subplot(132, projection='3d') +opt_coils.plot(ax=ax2, show=False,label='Standard optimized coils') +surface.plot(ax=ax2, show=False) +ax3 = fig.add_subplot(133, projection='3d') +opt_coils_alm.plot(ax=ax3, show=False,label='ALM optimized coils') +surface.plot(ax=ax3, show=False) +plt.legend() +plt.tight_layout() +plt.show() + +if EXPORT: + output_filepath = os.path.join(os.path.dirname(__file__), "output") + + """ Save the coils to a json file """ + init_coils.to_json(os.path.join(output_filepath, "init_coils_vmec_surface.json")) + opt_coils.to_json(os.path.join(output_filepath, "opt_coils_vmec_surface.json")) + + """ Save results in vtk format to analyze in Paraview """ + surface.to_vtk(os.path.join(output_filepath, "init_surface_vmec_surface.json"), field=init_field) + surface.to_vtk(os.path.join(output_filepath, "final_surface_vmec_surface.json"), field=opt_field) + init_coils.to_vtk(os.path.join(output_filepath, "init_coils_vmec_surface.json")) + opt_coils.to_vtk(os.path.join(output_filepath, "opt_coils_vmec_surface.json")) + opt_coils_alm.to_vtk(os.path.join(output_filepath, "opt_coils_alm_vmec_surface.json")) diff --git a/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_stochastic.py b/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_stochastic.py new file mode 100644 index 00000000..b2892200 --- /dev/null +++ b/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_stochastic.py @@ -0,0 +1,182 @@ +import os +number_of_processors_to_use = 1 # Parallelization, this should divide ntheta*nphi +os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' +from time import time + +import jax +import jax.numpy as jnp +import matplotlib.pyplot as plt + +import essos.augmented_lagrangian as alm +from essos.coil_perturbation import ( + GaussianSampler, + perturb_curves_statistic, + perturb_curves_systematic, +) +from essos.coils import Coils, CreateEquallySpacedCurves +from essos.fields import BiotSavart, Vmec +from essos.losses import custom_loss +from essos.surfaces import BdotN_over_B + + +""" Creating stochastic field losses """ +def copy_coils_from_field(field): + return field.coils.copy() + + +def perturbed_field_from_field(field, key, sampler): + coils = copy_coils_from_field(field) + base_key = jax.random.key(key) + split_keys = jax.random.split(base_key, 2) + coils = perturb_curves_systematic(coils, sampler, key=split_keys[0]) + coils = perturb_curves_statistic(coils, sampler, key=split_keys[1]) + return BiotSavart(coils) + + +def loss_bdotn_stochastic(field, surface, sampler, keys): + def perturbed_loss(key): + perturbed_field = perturbed_field_from_field(field, key, sampler) + bdotn_over_b = BdotN_over_B(surface, perturbed_field) + return jnp.sum(jnp.abs(bdotn_over_b)) + return jnp.mean(jax.vmap(perturbed_loss)(keys)) + + +def constraint_bdotn_stochastic(field, surface, sampler, keys, target_tol=1.0e-6): + def perturbed_square(key): + perturbed_field = perturbed_field_from_field(field, key, sampler) + return jnp.square(BdotN_over_B(surface, perturbed_field)) + expected_square = jnp.mean(jax.vmap(perturbed_square)(keys), axis=0) + return jnp.sqrt(jnp.sum(jnp.maximum(expected_square - target_tol, 0.0))) + + +def loss_length_constraint(field, max_coil_length): + return jnp.square(field.coils.length / max_coil_length - 1.0) + + +def loss_curvature_constraint(field, max_coil_curvature): + pointwise_curvature_loss = jnp.square(jnp.maximum(field.coils.curvature - max_coil_curvature, 0.0)) + return jnp.mean(pointwise_curvature_loss * jnp.linalg.norm(field.coils.gamma_dash, axis=-1), axis=1) + + +# Optimization parameters +maximum_function_evaluations = 10 + +MAX_COIL_LENGTH = 40.0 +MAX_COIL_CURVATURE = 0.5 +BDOTN_TARGET_TOL = 1.0e-6 +FOURIER_ORDER = 6 +N_SEGMENTS = FOURIER_ORDER * 10 +N_COILS = 4 +NTHETA = 32 +NPHI = 32 + +input_filepath = os.path.join(os.path.dirname(__file__), "..", "input_files") +vmec_input = os.path.join(input_filepath, "wout_LandremanPaul2021_QA_reactorScale_lowres.nc") + +""" Creating starting coils and surface """ +vmec = Vmec(vmec_input, ntheta=NTHETA, nphi=NPHI, range_torus="full torus") +surface = vmec.surface + +COIL_CURRENT = 1.0 +number_of_field_periods = vmec.nfp +major_radius_coils = vmec.r_axis +minor_radius_coils = vmec.r_axis / 1.5 +curves = CreateEquallySpacedCurves(n_curves=N_COILS,order=FOURIER_ORDER, R=major_radius_coils, + r=minor_radius_coils,n_segments=N_SEGMENTS, nfp=number_of_field_periods,stellsym=True) +coils_initial = Coils(curves=curves,currents=[COIL_CURRENT] * N_COILS) +field_initial = BiotSavart(coils_initial) + +""" Setting the stochastic sampling parameters """ +SIGMA = 0.01 +LENGTH_SCALE = 0.4 * jnp.pi +N_DERIVS = 2 +N_samples = 10 +sampler = GaussianSampler(coils_initial.curves.quadpoints, sigma=SIGMA, length_scale=LENGTH_SCALE, n_derivs=N_DERIVS) +stochastic_keys = jnp.arange(N_samples) + +""" Defining custom losses """ +L_normal_field = custom_loss(loss_bdotn_stochastic, "field", surface=surface, sampler=sampler, keys=stochastic_keys) +L_normal_field_constraint = custom_loss(constraint_bdotn_stochastic, "field", surface=surface, sampler=sampler, keys=stochastic_keys, target_tol=BDOTN_TARGET_TOL) +L_length_constraint = custom_loss(loss_length_constraint, "field", max_coil_length=MAX_COIL_LENGTH) +L_curvature_constraint = custom_loss(loss_curvature_constraint, "field", max_coil_curvature=MAX_COIL_CURVATURE) + +""" Defining total loss + setting dependencies """ +L_normal_field.dependencies = {"field": field_initial} +L_normal_field_constraint.dependencies = {"field": field_initial} +L_length_constraint.dependencies = {"field": field_initial} +L_curvature_constraint.dependencies = {"field": field_initial} + +""" Creating the constraints """ +penalty = 0.1 +multiplier = 0.5 +sq_grad = 0.0 +model_lagrangian = "Standard" + +beta = 2.0 +mu_max = 1.0e4 +alpha = 0.99 +gamma = 1.0e-2 +epsilon = 1.0e-8 +omega_tol = 1.0e-7 +eta_tol = 1.0e-7 + +normal_field_constraint = alm.eq(constraint_bdotn_stochastic, model_lagrangian=model_lagrangian, multiplier=multiplier, penalty=penalty, sq_grad=sq_grad) +length_constraint = alm.eq(loss_length_constraint, model_lagrangian=model_lagrangian, multiplier=multiplier, penalty=penalty, sq_grad=sq_grad) +curvature_constraint = alm.eq(loss_curvature_constraint, model_lagrangian=model_lagrangian, multiplier=multiplier, penalty=penalty, sq_grad=sq_grad) + +C_normal_field_constraint = alm.SelectiveConstraint(normal_field_constraint, "field", surface=surface, sampler=sampler, keys=stochastic_keys, target_tol=BDOTN_TARGET_TOL) +C_length_constraint = alm.SelectiveConstraint(length_constraint, "field", max_coil_length=MAX_COIL_LENGTH) +C_curvature_constraint = alm.SelectiveConstraint(curvature_constraint, "field", max_coil_curvature=MAX_COIL_CURVATURE) +C_total_constraint = alm.combine(C_normal_field_constraint, C_length_constraint, C_curvature_constraint) + +C_normal_field_constraint.dependencies = {"field": field_initial} +C_length_constraint.dependencies = {"field": field_initial} +C_curvature_constraint.dependencies = {"field": field_initial} +C_total_constraint.dependencies = {"field": field_initial} + +ALM = alm.ALM_model_jaxopt_lbfgsb(constraints=C_total_constraint, model_lagrangian=model_lagrangian, beta=beta, mu_max=mu_max, alpha=alpha, gamma=gamma, epsilon=epsilon, eta_tol=eta_tol, omega_tol=omega_tol) + +""" Optimizing with alm """ +lagrange_params = C_total_constraint.init(field_initial.dofs) +params = field_initial.dofs, lagrange_params +lag_state, grad, info = ALM.init(params) + +print(f"Optimizing coils with {maximum_function_evaluations} function evaluations using stochastic ALM.") +time0 = time() + +i = 0 +while i <= maximum_function_evaluations and (jnp.linalg.norm(grad[0]) > omega_tol or alm.norm_constraints(info[2]) > eta_tol): + params, lag_state, grad, info = ALM.update(params, lag_state, grad, info) + print(f"i: {i}, loss f: {info[0]:g}, loss L: {info[1]:g}, " f"infeasibility: {alm.total_infeasibility(info[2]):g}") + i += 1 + +field_optimized = C_normal_field_constraint.dofs_to_pytree(params[0])[0] +coils_optimized = field_optimized.coils + +print(f"Stochastic optimization with ALM took {time() - time0:.2f} seconds") + +BdotN_over_B_initial = BdotN_over_B(surface, field_initial) +BdotN_over_B_optimized = BdotN_over_B(surface, field_optimized) +curvature = jnp.mean(field_optimized.coils.curvature, axis=1) +length = jnp.max(jnp.ravel(field_optimized.coils.length)) +stochastic_loss_initial = L_normal_field(L_normal_field.starting_dofs) +stochastic_loss_final = L_normal_field(params[0]) + +print("Mean curvature:", curvature) +print("Length:", length) +print(f"Stochastic |BdotN/B| loss before optimization: {stochastic_loss_initial:.2e}") +print(f"Stochastic |BdotN/B| loss after optimization: {stochastic_loss_final:.2e}") +print(f"Maximum BdotN/B before optimization: {jnp.max(BdotN_over_B_initial):.2e}") +print(f"Maximum BdotN/B after optimization: {jnp.max(BdotN_over_B_optimized):.2e}") +print(f"Average BdotN/B before optimization: {jnp.average(jnp.abs(BdotN_over_B_initial)):.2e}") +print(f"Average BdotN/B after optimization: {jnp.average(jnp.abs(BdotN_over_B_optimized)):.2e}") + +fig = plt.figure(figsize=(8, 4)) +ax1 = fig.add_subplot(121, projection="3d") +ax2 = fig.add_subplot(122, projection="3d") +coils_initial.plot(ax=ax1, show=False) +surface.plot(ax=ax1, show=False) +coils_optimized.plot(ax=ax2, show=False) +surface.plot(ax=ax2, show=False) +plt.tight_layout() +plt.show() diff --git a/examples/coil_optimization/optimize_coils_vmec_surface_distance_forces.py b/examples/coil_optimization/optimize_coils_vmec_surface_distance_forces.py new file mode 100644 index 00000000..d8266e10 --- /dev/null +++ b/examples/coil_optimization/optimize_coils_vmec_surface_distance_forces.py @@ -0,0 +1,156 @@ +import os +from time import time +import jax.numpy as jnp +import matplotlib.pyplot as plt + +from essos.coils import Coils, CreateEquallySpacedCurves +from essos.fields import BiotSavart +from essos.surfaces import SurfaceRZFourier, BdotN_over_B +from essos.losses import custom_loss +from essos.objective_functions import ( + loss_coil_separation, + loss_coil_surface_distance, + loss_lorentz_force_coils, + loss_linkingnumber, +) + +# In this exmple, `scipy.optimize.least_squares` is used, but any other optimizer, e.g. from +# `scipy.optimize.minimize` or `jaxopt`, can be used as well and may even be preferable. +from scipy.optimize import least_squares + +input_filepath = os.path.join(os.path.dirname(__file__), "..", "input_files") +vmec_input = os.path.join(input_filepath, "wout_LandremanPaul2021_QA_reactorScale_lowres.nc") + +""" Creating starting coils and surface """ +N_COILS = 3; FOURIER_ORDER = 3; LARGE_R = 10; SMALL_R = 5.6; NFP = 2; N_SEGMENTS = 45; STELLSYM = True +COIL_CURRENT = 1. + +init_curves = CreateEquallySpacedCurves( + N_COILS, FOURIER_ORDER, LARGE_R, SMALL_R, n_segments=N_SEGMENTS, nfp=NFP, stellsym=STELLSYM +) +init_coils = Coils(curves=init_curves, currents=[COIL_CURRENT] * N_COILS) +init_field = BiotSavart(init_coils) +surface = SurfaceRZFourier.from_wout_file(vmec_input, s=1, ntheta=30, nphi=30, range_torus="half period") + +""" Setting the losses weights and targets """ +NORMAL_FIELD_WEIGHT = 1. +SURFACE_DISTANCE_WEIGHT = 1. +COIL_DISTANCE_WEIGHT = 1. +FORCE_WEIGHT = 1. +LINKING_WEIGHT = 1. + +MIN_SURFACE_DISTANCE = 0.2 +MIN_COIL_DISTANCE = 0.2 +FORCE_THRESHOLD = 0.5e6 +FORCE_POWER = 1 +BLOCK_SIZE = None + +""" Creating the loss functions """ +def loss(field, surface): + return jnp.sum(jnp.abs(BdotN_over_B(surface, field))) + + +def loss_surface_distance(field): + return loss_coil_surface_distance( + field.coils, surface, min_distance=MIN_SURFACE_DISTANCE, block_size=BLOCK_SIZE + ) + + +def loss_coil_distance(field): + return loss_coil_separation( + field.coils, min_separation=MIN_COIL_DISTANCE, block_size=BLOCK_SIZE + ) + + +def loss_force(field): + return loss_lorentz_force_coils( + field.coils, p=FORCE_POWER, threshold=FORCE_THRESHOLD, block_size=BLOCK_SIZE + ) + + +def loss_linking(field): + return loss_linkingnumber(field.coils, block_size=BLOCK_SIZE) + + +def print_loss_costs(label, field): + print(f"\n{label}") + print("normal_field:", loss(field, surface)) + print("coil_surface_distance:", loss_surface_distance(field)) + print("coil_coil_distance:", loss_coil_distance(field)) + print("lorentz_force:", loss_force(field)) + print("linking_number:", loss_linking(field)) + print( + "total_weighted:", + NORMAL_FIELD_WEIGHT * loss(field, surface) + + SURFACE_DISTANCE_WEIGHT * loss_surface_distance(field) + + COIL_DISTANCE_WEIGHT * loss_coil_distance(field) + + FORCE_WEIGHT * loss_force(field) + + LINKING_WEIGHT * loss_linking(field), + ) + + +""" Defining custom losses """ +L_normal_field = custom_loss(loss, "field", surface=surface) +L_surface_distance = custom_loss(loss_surface_distance, "field") +L_coil_distance = custom_loss(loss_coil_distance, "field") +L_force = custom_loss(loss_force, "field") +L_linking = custom_loss(loss_linking, "field") + +""" Defining total loss + setting dependencies """ +L_total = ( + NORMAL_FIELD_WEIGHT * L_normal_field + + SURFACE_DISTANCE_WEIGHT * L_surface_distance + + COIL_DISTANCE_WEIGHT * L_coil_distance + + FORCE_WEIGHT * L_force + + LINKING_WEIGHT * L_linking +) +L_total.dependencies = {"field": init_field} + +""" Optimizing the total loss """ +t_start = time() +res = least_squares( + L_total, + L_total.starting_dofs, + L_total.grad, + verbose=2, + ftol=1e-5, + gtol=1e-5, + xtol=1e-14, + max_nfev=200, +) +t_end = time() + +print(f"\nOptimization took {t_end - t_start:.2f} seconds") +print("Initial loss:", L_total(L_total.starting_dofs)) +print("Loss after optimization:", L_total(res.x)) + +opt_field = L_total.dofs_to_pytree(res.x)["field"] +opt_coils = opt_field.coils + +print_loss_costs("Loss costs before optimization", init_field) +print_loss_costs("Loss costs after optimization", opt_field) + +fig = plt.figure(figsize=(8, 4)) + +ax1 = fig.add_subplot(121, projection="3d") +init_coils.plot(ax=ax1, show=False) +surface.plot(ax=ax1, show=False) +ax2 = fig.add_subplot(122, projection="3d") +opt_coils.plot(ax=ax2, show=False) +surface.plot(ax=ax2, show=False) +plt.tight_layout() +plt.show() + +EXPORT = False +if EXPORT: + output_filepath = os.path.join(os.path.dirname(__file__), "output") + + """ Save the coils to a json file """ + init_coils.to_json(os.path.join(output_filepath, "init_coils_vmec_surface_distance_forces.json")) + opt_coils.to_json(os.path.join(output_filepath, "opt_coils_vmec_surface_distance_forces.json")) + + """ Save results in vtk format to analyze in Paraview """ + surface.to_vtk(os.path.join(output_filepath, "init_surface_vmec_surface_distance_forces.json"), field=init_field) + surface.to_vtk(os.path.join(output_filepath, "final_surface_vmec_surface_distance_forces.json"), field=opt_field) + init_coils.to_vtk(os.path.join(output_filepath, "init_coils_vmec_surface_distance_forces.json")) + opt_coils.to_vtk(os.path.join(output_filepath, "opt_coils_vmec_surface_distance_forces.json")) diff --git a/examples/optimize_multiple_objectives.py b/examples/coil_optimization/optimize_multiple_objectives.py similarity index 90% rename from examples/optimize_multiple_objectives.py rename to examples/coil_optimization/optimize_multiple_objectives.py index b6cf47b6..e826a805 100644 --- a/examples/optimize_multiple_objectives.py +++ b/examples/coil_optimization/optimize_multiple_objectives.py @@ -1,4 +1,5 @@ +import os import jax.numpy as jnp from essos.fields import BiotSavart from essos.fields import Vmec @@ -6,7 +7,7 @@ from essos.objective_functions import loss_normB_axis,loss_bdotn_over_b,loss_coil_length, loss_coil_curvature, loss_BdotN from essos.multiobjectiveoptimizer import MultiObjectiveOptimizer -vmec = Vmec("./input_files/wout_LandremanPaul2021_QA_reactorScale_lowres.nc", ntheta=32, nphi=32, range_torus='half period') +vmec = Vmec(os.path.join(os.path.dirname(__file__), "..", "input_files", "wout_LandremanPaul2021_QA_reactorScale_lowres.nc"), ntheta=32, nphi=32, range_torus='half period') # inputs manager = MultiObjectiveOptimizer( diff --git a/examples/compare_guidingcenter_fullorbit.py b/examples/compare_guidingcenter_fullorbit.py deleted file mode 100644 index 27a3b518..00000000 --- a/examples/compare_guidingcenter_fullorbit.py +++ /dev/null @@ -1,91 +0,0 @@ -import os -number_of_processors_to_use = 1 # Parallelization, this should divide nparticles -os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' -from jax import vmap -from time import time -import jax.numpy as jnp -import matplotlib.pyplot as plt -from essos.fields import BiotSavart -from essos.coils import Coils_from_json -from essos.constants import PROTON_MASS, ONE_EV -from essos.dynamics import Tracing, Particles -from jax import block_until_ready - -# Input parameters -tmax = 1.e-4 -dt_fo=1.e-9 -nparticles_per_core=2 -nparticles = number_of_processors_to_use*nparticles_per_core -R0 = jnp.linspace(1.23, 1.27, nparticles) -trace_tolerance = 1e-5 -num_steps_gc = 5000 -num_steps_fo = int(tmax/dt_fo) -mass=PROTON_MASS -energy=5000*ONE_EV - -# Load coils and field -json_file = os.path.join(os.path.dirname(__file__), 'input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') -coils = Coils_from_json(json_file) -field = BiotSavart(coils) - -# Initialize particles -Z0 = jnp.zeros(nparticles) -phi0 = jnp.zeros(nparticles) -initial_xyz=jnp.array([R0*jnp.cos(phi0), R0*jnp.sin(phi0), Z0]).T -initial_vparallel_over_v = jnp.linspace(-0.1, 0.1, nparticles) -particles = Particles(initial_xyz=initial_xyz, mass=mass, energy=energy, field=field, initial_vparallel_over_v=initial_vparallel_over_v) - -# Trace in ESSOS -time0 = time() -tracing_guidingcenter = Tracing(field=field, model='GuidingCenterAdaptative', particles=particles, - maxtime=tmax,times_to_trace=num_steps_gc, atol=trace_tolerance,rtol=trace_tolerance) -trajectories_guidingcenter = block_until_ready(tracing_guidingcenter.trajectories) -print(f"ESSOS guiding center tracing took {time()-time0:.2f} seconds") - -time0 = time() -tracing_fullorbit = Tracing(field=field, model='FullOrbit_Boris', particles=particles, - maxtime=tmax, times_to_trace=num_steps_fo,timestep=dt_fo) -trajectories_fullorbit = block_until_ready(tracing_fullorbit.trajectories) -print(f"ESSOS full orbit tracing took {time()-time0:.2f} seconds") - -# Plot trajectories, velocity parallel to the magnetic field, and energy error -fig = plt.figure(figsize=(9, 8)) -ax1 = fig.add_subplot(221, projection='3d') -ax2 = fig.add_subplot(222) -ax3 = fig.add_subplot(223) -ax4 = fig.add_subplot(224) - -coils.plot(ax=ax1, show=False) -tracing_guidingcenter.plot(ax=ax1, show=False) -tracing_fullorbit.plot(ax=ax1, show=False) - -for i, (trajectory_gc, trajectory_fo) in enumerate(zip(trajectories_guidingcenter, trajectories_fullorbit)): - ax2.plot(tracing_guidingcenter.times, jnp.abs(tracing_guidingcenter.energy[i]-particles.energy)/particles.energy, '-', label=f'Particle {i+1} GC', linewidth=1.0, alpha=0.7) - ax2.plot(tracing_fullorbit.times, jnp.abs(tracing_fullorbit.energy[i]-particles.energy)/particles.energy, '--', label=f'Particle {i+1} FO', linewidth=1.0, markersize=0.5, alpha=0.7) - def compute_v_parallel(trajectory_t): - magnetic_field_unit_vector = field.B(trajectory_t[:3]) / field.AbsB(trajectory_t[:3]) - return jnp.dot(trajectory_t[3:], magnetic_field_unit_vector) - v_parallel_fo = vmap(compute_v_parallel)(trajectory_fo) - ax3.plot(tracing_guidingcenter.times, trajectory_gc[:, 3] / particles.total_speed, '-', label=f'Particle {i+1} GC', linewidth=1.1, alpha=0.95) - ax3.plot(tracing_fullorbit.times, v_parallel_fo / particles.total_speed, '--', label=f'Particle {i+1} FO', linewidth=0.5, markersize=0.5, alpha=0.2) - # ax4.plot(jnp.sqrt(trajectory_gc[:,0]**2+trajectory_gc[:,1]**2), trajectory_gc[:, 2], '-', label=f'Particle {i+1} GC', linewidth=1.5, alpha=0.3) - # ax4.plot(jnp.sqrt(trajectory_fo[:,0]**2+trajectory_fo[:,1]**2), trajectory_fo[:, 2], '--', label=f'Particle {i+1} FO', linewidth=1.5, markersize=0.5, alpha=0.2) -tracing_guidingcenter.poincare_plot(ax=ax4, show=False, color='k', label=f'GC', shifts=[jnp.pi/2])#, 0]) -tracing_fullorbit.poincare_plot( ax=ax4, show=False, color='r', label=f'FO', shifts=[jnp.pi/2])#, 0]) - -ax2.set_xlabel('Time (s)') -ax2.set_ylabel('Relative Energy Error') -ax3.set_ylabel(r'$v_{\parallel}/v$') -ax2.legend(loc='upper right') -ax3.set_xlabel('Time (s)') -ax3.legend(loc='upper right') -ax4.set_xlabel('R (m)') -ax4.set_ylabel('Z (m)') -ax4.legend(loc='upper right') -plt.tight_layout() -plt.show() - - -## Save results in vtk format to analyze in Paraview -# tracing.to_vtk('trajectories') -# coils.to_vtk('coils') \ No newline at end of file diff --git a/examples/comparisons_SIMSOPT/coils_biotsavart_SIMSOPT_vs_ESSOS.py b/examples/comparisons_SIMSOPT/coils_biotsavart_SIMSOPT_vs_ESSOS.py deleted file mode 100644 index c89fc99f..00000000 --- a/examples/comparisons_SIMSOPT/coils_biotsavart_SIMSOPT_vs_ESSOS.py +++ /dev/null @@ -1,260 +0,0 @@ -import os -from time import time -import jax.numpy as jnp -import matplotlib.pyplot as plt -from jax import block_until_ready -from essos.fields import BiotSavart as BiotSavart_essos -from essos.coils import Coils_from_simsopt, Curves_from_simsopt -from simsopt import load -from simsopt.geo import CurveXYZFourier, curves_to_vtk -from simsopt.field import BiotSavart as BiotSavart_simsopt, coils_via_symmetries -from simsopt.configs import get_ncsx_data, get_w7x_data, get_hsx_data, get_giuliani_data - -output_dir = os.path.join(os.path.dirname(__file__), 'output') -if not os.path.exists(output_dir): - os.makedirs(output_dir) - -list_segments = [30, 100, 300, 1000, 3000] - -LandremanPaulQA_json_file = os.path.join(os.path.dirname(__file__), '..', 'input_files', 'SIMSOPT_biot_savart_LandremanPaulQA.json') -nfp_array = [3, 2, 5, 4, 2] -curves_array = [get_ncsx_data()[0], LandremanPaulQA_json_file, get_w7x_data()[0], get_hsx_data()[0], get_giuliani_data()[0]] -currents_array = [get_ncsx_data()[1], None, get_w7x_data()[1], get_hsx_data()[1], get_giuliani_data()[1]] -name_array = ["NCSX", "QA(json)", "W7-X", "HSX", "Giuliani"] - -print(f'Output being saved to {output_dir}') -print(f'SIMSOPT LandremanPaulQA json file location: {LandremanPaulQA_json_file}') -for nfp, curves_stel, currents_stel, name in zip(nfp_array, curves_array, currents_array, name_array): - print(f' Running {name} and saving to output directory...') - if currents_stel is None: - json_file_stel = curves_stel - field_simsopt = load(json_file_stel) - coils_simsopt = field_simsopt.coils - curves_simsopt = [coil.curve for coil in coils_simsopt] - currents_simsopt = [coil.current for coil in coils_simsopt] - coils_essos = Coils_from_simsopt(json_file_stel, nfp) - curves_essos = Curves_from_simsopt(json_file_stel, nfp) - else: - coils_simsopt = coils_via_symmetries(curves_stel, currents_stel, nfp, True) - curves_simsopt = [c.curve for c in coils_simsopt] - currents_simsopt = [c.current for c in coils_simsopt] - field_simsopt = BiotSavart_simsopt(coils_simsopt) - - coils_essos = Coils_from_simsopt(coils_simsopt, nfp) - curves_essos = Curves_from_simsopt(curves_simsopt, nfp) - - field_essos = BiotSavart_essos(coils_essos) - - coils_essos_to_simsopt = coils_essos.to_simsopt() - curves_essos_to_simsopt = curves_essos.to_simsopt() - field_essos_to_simsopt = BiotSavart_simsopt(coils_essos_to_simsopt) - - curves_to_vtk(curves_simsopt, os.path.join(output_dir,f"curves_simsopt_{name}")) - curves_essos.to_vtk(os.path.join(output_dir,f"curves_essos_{name}")) - curves_to_vtk(curves_essos_to_simsopt, os.path.join(output_dir,f"curves_essos_to_simsopt_{name}")) - - base_coils_simsopt = coils_simsopt[:int(len(coils_simsopt)/2/nfp)] - R = jnp.mean(jnp.array([jnp.sqrt(coil.curve.x[coil.curve.local_dof_names.index('xc(0)')]**2 - +coil.curve.x[coil.curve.local_dof_names.index('yc(0)')]**2) - for coil in base_coils_simsopt])) - x = jnp.array([R+0.01,R,R]) - y = jnp.array([R,R+0.01,R-0.01]) - z = jnp.array([0.05,0.06,0.07]) - - positions = jnp.array((x,y,z)) - - len_list_segments = len(list_segments) - t_gamma_avg_essos = jnp.zeros(len_list_segments) - t_gamma_avg_simsopt = jnp.zeros(len_list_segments) - gamma_error_avg = jnp.zeros(len_list_segments) - t_gammadash_avg_essos = jnp.zeros(len_list_segments) - t_gammadash_avg_simsopt = jnp.zeros(len_list_segments) - gammadash_error_avg = jnp.zeros(len_list_segments) - t_gammadashdash_avg_essos = jnp.zeros(len_list_segments) - t_gammadashdash_avg_simsopt = jnp.zeros(len_list_segments) - gammadashdash_error_avg = jnp.zeros(len_list_segments) - t_curvature_avg_essos = jnp.zeros(len_list_segments) - t_curvature_avg_simsopt = jnp.zeros(len_list_segments) - curvature_error_avg = jnp.zeros(len_list_segments) - t_B_avg_essos = jnp.zeros(len_list_segments) - t_B_avg_simsopt = jnp.zeros(len_list_segments) - B_error_avg = jnp.zeros(len_list_segments) - t_dB_by_dX_avg_essos = jnp.zeros(len_list_segments) - t_dB_by_dX_avg_simsopt = jnp.zeros(len_list_segments) - dB_by_dX_error_avg = jnp.zeros(len_list_segments) - - gamma_error_simsopt_to_essos = 0 - gamma_error_essos_to_simsopt = 0 - - for i, (coil_simsopt, coil_essos_gamma, coil_essos_to_simsopt) in enumerate(zip(coils_simsopt, coils_essos.gamma, coils_essos_to_simsopt)): - gamma_error_simsopt_to_essos += jnp.linalg.norm(coil_simsopt.curve.gamma()-coil_essos_gamma) - gamma_error_essos_to_simsopt += jnp.linalg.norm(coil_simsopt.curve.gamma()-coil_essos_to_simsopt.curve.gamma()) - - B_error_avg_simsopt_to_essos = 0 - B_error_avg_essos_to_simsopt = 0 - for j, position in enumerate(positions): - field_simsopt.set_points([position]) - field_essos_to_simsopt.set_points([position]) - B_simsopt = field_simsopt.B() - B_essos_to_simsopt = field_essos_to_simsopt.B() - B_simsopt_to_essos = field_essos.B(position) - B_error_avg_simsopt_to_essos += jnp.abs(jnp.linalg.norm(B_simsopt) - jnp.linalg.norm(B_simsopt_to_essos)) - B_error_avg_essos_to_simsopt += jnp.abs(jnp.linalg.norm(B_simsopt) - jnp.linalg.norm(B_essos_to_simsopt)) - B_error_avg_simsopt_to_essos = B_error_avg_simsopt_to_essos/len(positions) - B_error_avg_essos_to_simsopt = B_error_avg_essos_to_simsopt/len(positions) - - fig = plt.figure(figsize = (8, 6)) - X_axis = jnp.arange(2) - plt.bar(X_axis[0] - 0.2, gamma_error_simsopt_to_essos+1e-19, 0.3, label='SIMSOPT to ESSOS coils', color='blue', edgecolor='black', hatch='/') - plt.bar(X_axis[0] + 0.2, gamma_error_essos_to_simsopt+1e-19, 0.3, label='ESSOS to SIMSOPT coils', color='red', edgecolor='black', hatch='-') - plt.bar(X_axis[1] - 0.2, B_error_avg_simsopt_to_essos+1e-19, 0.3, label=r'SIMSOPT to ESSOS $B$', color='blue', edgecolor='black', hatch='||') - plt.bar(X_axis[1] + 0.2, B_error_avg_essos_to_simsopt+1e-19, 0.3, label=r'ESSOS to SIMSOPT $B$', color='red', edgecolor='black', hatch='*') - plt.xticks(X_axis, ['Coil Error', 'B Error']) - plt.xlabel('Parameter', fontsize=14) - plt.ylabel('Error Magnitude', fontsize=14) - plt.yscale('log') - plt.ylim(1e-20, 1e-11) - plt.legend(fontsize=14) - plt.grid(axis='y') - plt.title(f"{name}", fontsize=14) - plt.tight_layout() - plt.savefig(os.path.join(output_dir,f"error_gamma_B_SIMSOPT_vs_ESSOS_{name}.pdf"), transparent=True) - plt.close() - - def update_nsegments_simsopt(curve_simsopt, n_segments): - new_curve = CurveXYZFourier(n_segments, curve_simsopt.order) - new_curve.x = curve_simsopt.x - return new_curve - - for index, n_segments in enumerate(list_segments): - coils_essos.n_segments = n_segments - - base_curves_simsopt = [update_nsegments_simsopt(coil_simsopt.curve, n_segments) for coil_simsopt in base_coils_simsopt] - coils_simsopt = coils_via_symmetries(base_curves_simsopt, currents_simsopt[0:len(base_coils_simsopt)], nfp, True) - curves_simsopt = [c.curve for c in coils_simsopt] - - [curve.gamma() for curve in curves_simsopt] - coils_essos.gamma - - start_time = time() - gamma_curves_simsopt = block_until_ready(jnp.array([curve.gamma() for curve in curves_simsopt])) - t_gamma_avg_simsopt = t_gamma_avg_simsopt.at[index].set(t_gamma_avg_simsopt[index] + time() - start_time) - - start_time = time() - gamma_curves_essos = block_until_ready(jnp.array(coils_essos.gamma)) - t_gamma_avg_essos = t_gamma_avg_essos.at[index].set(t_gamma_avg_essos[index] + time() - start_time) - - start_time = time() - gammadash_curves_simsopt = block_until_ready(jnp.array([curve.gammadash() for curve in curves_simsopt])) - t_gammadash_avg_simsopt = t_gammadash_avg_simsopt.at[index].set(t_gammadash_avg_simsopt[index] + time() - start_time) - - start_time = time() - gammadash_curves_essos = block_until_ready(jnp.array(coils_essos.gamma_dash)) - t_gammadash_avg_essos = t_gammadash_avg_essos.at[index].set(t_gammadash_avg_essos[index] + time() - start_time) - - start_time = time() - gammadashdash_curves_simsopt = block_until_ready(jnp.array([curve.gammadashdash() for curve in curves_simsopt])) - t_gammadashdash_avg_simsopt = t_gammadashdash_avg_simsopt.at[index].set(t_gammadashdash_avg_simsopt[index] + time() - start_time) - - start_time = time() - gammadashdash_curves_essos = block_until_ready(jnp.array(coils_essos.gamma_dashdash)) - t_gammadashdash_avg_essos = t_gammadashdash_avg_essos.at[index].set(t_gammadashdash_avg_essos[index] + time() - start_time) - - start_time = time() - curvature_curves_simsopt = block_until_ready(jnp.array([curve.kappa() for curve in curves_simsopt])) - t_curvature_avg_simsopt = t_curvature_avg_simsopt.at[index].set(t_curvature_avg_simsopt[index] + time() - start_time) - - start_time = time() - curvature_curves_essos = block_until_ready(jnp.array(coils_essos.curvature)) - t_curvature_avg_essos = t_curvature_avg_essos.at[index].set(t_curvature_avg_essos[index] + time() - start_time) - - gamma_error_avg = gamma_error_avg. at[index].set(gamma_error_avg[index] + jnp.linalg.norm(gamma_curves_essos - gamma_curves_simsopt)) - gammadash_error_avg = gammadash_error_avg. at[index].set(gammadash_error_avg[index] + jnp.linalg.norm(gammadash_curves_essos - gammadash_curves_simsopt)) - gammadashdash_error_avg = gammadashdash_error_avg.at[index].set(gammadashdash_error_avg[index] + jnp.linalg.norm(gammadashdash_curves_essos - gammadashdash_curves_simsopt)) - curvature_error_avg = curvature_error_avg.at[index].set(curvature_error_avg[index] + jnp.linalg.norm(curvature_curves_essos - curvature_curves_simsopt)) - - field_essos = BiotSavart_essos(coils_essos) - field_simsopt = BiotSavart_simsopt(coils_simsopt) - - for j, position in enumerate(positions): - field_essos.B(position) - time1 = time() - result_B_essos = field_essos.B(position) - t_B_avg_essos = t_B_avg_essos.at[index].set(t_B_avg_essos[index] + time() - time1) - normB_essos = jnp.linalg.norm(result_B_essos) - - field_simsopt.set_points(jnp.array([position])) - field_simsopt.B() - time3 = time() - field_simsopt.set_points(jnp.array([position])) - result_simsopt = field_simsopt.B() - t_B_avg_simsopt = t_B_avg_simsopt.at[index].set(t_B_avg_simsopt[index] + time() - time3) - normB_simsopt = jnp.linalg.norm(jnp.array(result_simsopt)) - - B_error_avg = B_error_avg.at[index].set(B_error_avg[index] + jnp.abs(normB_essos - normB_simsopt)) - - field_essos.dB_by_dX(position) - time1 = time() - field_simsopt.set_points(jnp.array([position])) - result_dB_by_dX_essos = field_essos.dB_by_dX(position) - t_dB_by_dX_avg_essos = t_dB_by_dX_avg_essos.at[index].set(t_dB_by_dX_avg_essos[index] + time() - time1) - norm_dB_by_dX_essos = jnp.linalg.norm(result_dB_by_dX_essos) - - field_simsopt.dB_by_dX() - time3 = time() - field_simsopt.set_points(jnp.array([position])) - result_dB_by_dX_simsopt = field_simsopt.dB_by_dX() - t_dB_by_dX_avg_simsopt = t_dB_by_dX_avg_simsopt.at[index].set(t_dB_by_dX_avg_simsopt[index] + time() - time3) - norm_dB_by_dX_simsopt = jnp.linalg.norm(jnp.array(result_dB_by_dX_simsopt)) - - dB_by_dX_error_avg = dB_by_dX_error_avg.at[index].set(dB_by_dX_error_avg[index] + jnp.abs(norm_dB_by_dX_essos - norm_dB_by_dX_simsopt)) - - X_axis = jnp.arange(len_list_segments) - - fig = plt.figure(figsize = (8, 6)) - plt.bar(X_axis-0.2, B_error_avg, 0.1, label = r"$B_{\text{essos}} - B_{\text{simsopt}}$", color="green", edgecolor="black", hatch="/") - plt.bar(X_axis-0.1, dB_by_dX_error_avg, 0.1, label = r"${B'}_{\text{essos}} - {B'}_{\text{simsopt}}$", color="purple", edgecolor="black", hatch="x") - plt.bar(X_axis+0.0, gamma_error_avg, 0.1, label = r"$\Gamma_{\text{essos}} - \Gamma_{\text{simsopt}}$", color="orange", edgecolor="black", hatch="|") - plt.bar(X_axis+0.1, gammadash_error_avg, 0.1, label = r"${\Gamma'}_{\text{essos}} - {\Gamma'}_{\text{simsopt}}$", color="gray", edgecolor="black", hatch="-") - plt.bar(X_axis+0.2, gammadashdash_error_avg, 0.1, label = r"${\Gamma''}_{\text{essos}} - {\Gamma''}_{\text{simsopt}}$", color="black", edgecolor="black", hatch="*") - plt.bar(X_axis+0.3, curvature_error_avg, 0.1, label = r"$\kappa_{\text{essos}} - \kappa_{\text{simsopt}}$", color="brown", edgecolor="black", hatch="\\") - plt.xticks(X_axis, list_segments) - plt.xlabel("Number of segments of each coil", fontsize=14) - plt.ylabel(f"Difference SIMSOPT vs ESSOS", fontsize=14) - plt.tick_params(axis='both', which='major', labelsize=14) - plt.tick_params(axis='both', which='minor', labelsize=14) - plt.legend(fontsize=14) - plt.yscale("log") - plt.grid(axis='y') - plt.ylim(1e-18, 1e-10) - plt.title(f"{name}", fontsize=14) - plt.tight_layout() - plt.savefig(os.path.join(output_dir,f"error_BiotSavart_SIMSOPT_vs_ESSOS_{name}.pdf"), transparent=True) - plt.close() - - fig = plt.figure(figsize = (8, 6)) - plt.bar(X_axis - 0.30, t_B_avg_essos, 0.05, label = r'B ESSOS', color="red", edgecolor="black") - plt.bar(X_axis - 0.25, t_B_avg_simsopt, 0.05, label = r'B SIMSOPT', color="blue", edgecolor="black") - plt.bar(X_axis - 0.20, t_dB_by_dX_avg_essos, 0.05, label = r"$B'$ ESSOS", color="red", edgecolor="black") - plt.bar(X_axis - 0.15, t_dB_by_dX_avg_simsopt, 00.05, label = r"$B'$ SIMSOPT", color="blue", edgecolor="black") - plt.bar(X_axis - 0.10, t_gamma_avg_essos, 0.05, label = r'$\Gamma$ ESSOS', color="red", edgecolor="black", hatch="//") - plt.bar(X_axis - 0.05, t_gamma_avg_simsopt, 0.05, label = r'$\Gamma$ SIMSOPT', color="blue", edgecolor="black", hatch="-") - plt.bar(X_axis + 0.0, t_gammadash_avg_essos, 0.05, label = r"${\Gamma'}$ ESSOS", color="red", edgecolor="black", hatch="\\") - plt.bar(X_axis + 0.05, t_gammadash_avg_simsopt, 0.05, label = r"${\Gamma'}$ SIMSOPT", color="blue", edgecolor="black", hatch="||") - plt.bar(X_axis + 0.10, t_gammadashdash_avg_essos, 0.05, label = r"${\Gamma''}$ ESSOS", color="red", edgecolor="black", hatch="*") - plt.bar(X_axis + 0.15, t_gammadashdash_avg_simsopt, 0.05, label = r"${\Gamma''}$ SIMSOPT", color="blue", edgecolor="black", hatch="|") - plt.bar(X_axis + 0.20, t_curvature_avg_essos, 0.05, label = r"$\kappa$ ESSOS", color="red", edgecolor="black", hatch="x") - plt.bar(X_axis + 0.25, t_curvature_avg_simsopt, 0.05, label = r"$\kappa$ SIMSOPT", color="blue", edgecolor="black", hatch="+") - plt.tick_params(axis='both', which='major', labelsize=14) - plt.tick_params(axis='both', which='minor', labelsize=14) - plt.xticks(X_axis, list_segments) - plt.xlabel("Number of segments of each coil", fontsize=14) - plt.ylabel("Time to evaluate SIMSOPT vs ESSOS (s)", fontsize=14) - plt.grid(axis='y') - # plt.gca().set_ylim((None,0.03)) - plt.yscale("log") - plt.legend(fontsize=14) - plt.title(f"{name}", fontsize=14) - plt.tight_layout() - plt.savefig(os.path.join(output_dir,f"time_BiotSavart_SIMSOPT_vs_ESSOS_{name}.pdf"), transparent=True) - plt.close() diff --git a/examples/comparisons_SIMSOPT/fieldlines_SIMSOPT_vs_ESSOS.py b/examples/comparisons_SIMSOPT/fieldlines_SIMSOPT_vs_ESSOS.py deleted file mode 100644 index fd0473e9..00000000 --- a/examples/comparisons_SIMSOPT/fieldlines_SIMSOPT_vs_ESSOS.py +++ /dev/null @@ -1,195 +0,0 @@ -import os -import time -import jax.numpy as jnp -from jax import block_until_ready -from simsopt import load -from simsopt.field import (particles_to_vtk, compute_fieldlines, plot_poincare_data) -from essos.coils import Coils_from_simsopt -from essos.dynamics import Tracing -from essos.fields import BiotSavart as BiotSavart_essos -import matplotlib.pyplot as plt - -tmax_fl = 150 -nfieldlines = 3 -axis_shft=0.02 -R0 = jnp.linspace(1.2125346+axis_shft, 1.295-axis_shft, nfieldlines) -nfp = 2 -trace_tolerance_SIMSOPT_array = [1e-5, 1e-7, 1e-9, 1e-11, 1e-13] -trace_tolerance_ESSOS = 1e-7 - -Z0 = jnp.zeros(nfieldlines) -phi0 = jnp.zeros(nfieldlines) - -phis_poincare = [(i/4)*(2*jnp.pi/nfp) for i in range(4)] - -output_dir = os.path.join(os.path.dirname(__file__), 'output') -if not os.path.exists(output_dir): - os.makedirs(output_dir) - -LandremanPaulQA_json_file = os.path.join(os.path.dirname(__file__), '..', 'input_files', 'SIMSOPT_biot_savart_LandremanPaulQA.json') -field_simsopt = load(LandremanPaulQA_json_file) -field_essos = BiotSavart_essos(Coils_from_simsopt(LandremanPaulQA_json_file, nfp)) - -fieldlines_SIMSOPT_array = [] -time_SIMSOPT_array = [] -avg_steps_SIMSOPT = 0 - -print(f'Output being saved to {output_dir}') -print(f'SIMSOPT LandremanPaulQA json file location: {LandremanPaulQA_json_file}') -for trace_tolerance_SIMSOPT in trace_tolerance_SIMSOPT_array: - print(f' Tracing SIMSOPT fieldlines with tolerance={trace_tolerance_SIMSOPT}') - t1 = time.time() - fieldlines_SIMSOPT_this_tolerance, fieldlines_SIMSOPT_phi_hits = block_until_ready(compute_fieldlines(field_simsopt, R0, Z0, tmax=tmax_fl, tol=trace_tolerance_SIMSOPT, phis=phis_poincare)) - time_SIMSOPT_array.append(time.time()-t1) - avg_steps_SIMSOPT += sum([len(l) for l in fieldlines_SIMSOPT_this_tolerance])//nfieldlines - print(f" Time for SIMSOPT tracing={time.time()-t1:.3f}s. Avg num steps={avg_steps_SIMSOPT}") - fieldlines_SIMSOPT_array.append(fieldlines_SIMSOPT_this_tolerance) - -particles_to_vtk(fieldlines_SIMSOPT_this_tolerance, os.path.join(output_dir,f'fieldlines_SIMSOPT')) -# plot_poincare_data(fieldlines_phi_hits, phis_poincare, os.path.join(output_dir,f'poincare_fieldline_SIMSOPT.pdf'), dpi=150) - -# Trace in ESSOS -num_steps_essos = int(jnp.mean(jnp.array([len(fieldlines_SIMSOPT[0]) for fieldlines_SIMSOPT in fieldlines_SIMSOPT_array]))) -time_essos = jnp.linspace(0, tmax_fl, num_steps_essos) - -print(f'Tracing ESSOS fieldlines with tolerance={trace_tolerance_ESSOS}') -t1 = time.time() -tracing = block_until_ready(Tracing(field=field_essos, model='FieldLine', initial_conditions=jnp.array([R0*jnp.cos(phi0), R0*jnp.sin(phi0), Z0]).T, - maxtime=tmax_fl, timesteps=num_steps_essos, tol_step_size=trace_tolerance_ESSOS)) -fieldlines_ESSOS = tracing.trajectories -time_ESSOS = time.time()-t1 -print(f" Time for ESSOS tracing={time.time()-t1:.3f}s. Num steps={len(fieldlines_ESSOS[0])}") - -tracing.to_vtk(os.path.join(output_dir,f'fieldlines_ESSOS')) -# tracing.poincare_plot(phis_poincare, show=False) - -print('Plotting the results to output directory...') -# Plot time comparison in a bar chart -labels = [f'SIMSOPT\nTol={tol}' for tol in trace_tolerance_SIMSOPT_array] + [f'ESSOS\nTol={trace_tolerance_ESSOS}'] -times = time_SIMSOPT_array + [time_ESSOS] -plt.figure() -bars = plt.bar(labels, times, color=['blue']*len(trace_tolerance_SIMSOPT_array) + ['red'], edgecolor=['black']*len(trace_tolerance_SIMSOPT_array) + ['black'], hatch=['//']*len(trace_tolerance_SIMSOPT_array) + ['|']) -plt.xlabel('Tracing Tolerance of SIMSOPT') -plt.ylabel('Time (s)') -plt.xticks(rotation=45) -plt.tight_layout() -blue_patch = plt.Line2D([0], [0], color='blue', lw=4, label='SIMSOPT', linestyle='--') -orange_patch = plt.Line2D([0], [0], color='red', lw=4, label=f'ESSOS', linestyle='-') -plt.legend(handles=[blue_patch, orange_patch]) -plt.savefig(os.path.join(output_dir, 'times_fieldlines_SIMSOPT_vs_ESSOS.pdf'), dpi=150) -plt.close() - -def interpolate_ESSOS_to_SIMSOPT(fieldine_SIMSOPT, fieldline_ESSOS): - time_SIMSOPT = jnp.array(fieldine_SIMSOPT)[:, 0] # Time values from fieldlines_SIMSOPT - # coords_SIMSOPT = jnp.array(fieldine_SIMSOPT)[:, 1:] # Coordinates (x, y, z) from fieldlines_SIMSOPT - coords_ESSOS = jnp.array(fieldline_ESSOS) - - interp_x = jnp.interp(time_SIMSOPT, time_essos, coords_ESSOS[:, 0]) - interp_y = jnp.interp(time_SIMSOPT, time_essos, coords_ESSOS[:, 1]) - interp_z = jnp.interp(time_SIMSOPT, time_essos, coords_ESSOS[:, 2]) - - coords_ESSOS_interp = jnp.column_stack([ interp_x, interp_y, interp_z]) - - return coords_ESSOS_interp - -relative_error_array = [] -for i, fieldlines_SIMSOPT in enumerate(fieldlines_SIMSOPT_array): - fieldlines_ESSOS_interp = [interpolate_ESSOS_to_SIMSOPT(fieldlines_SIMSOPT[i], fieldlines_ESSOS[i]) for i in range(nfieldlines)] - tracing.trajectories = fieldlines_ESSOS_interp - if i==len(trace_tolerance_SIMSOPT_array)-1: tracing.to_vtk(os.path.join(output_dir,f'fieldlines_ESSOS_interp')) - - relative_error_fieldlines_SIMSOPT_vs_ESSOS = [] - plt.figure() - for j in range(nfieldlines): - this_fieldline_SIMSOPT = jnp.array(fieldlines_SIMSOPT[j])[:,1:] - this_fieldlines_ESSOS = fieldlines_ESSOS_interp[j] - average_relative_error = [] - for fieldline_SIMSOPT_t, fieldline_ESSOS_t in zip(this_fieldline_SIMSOPT, this_fieldlines_ESSOS): - relative_error_x = jnp.abs(fieldline_SIMSOPT_t[0] - fieldline_ESSOS_t[0])/(jnp.abs(fieldline_SIMSOPT_t[0])+1e-12) - relative_error_y = jnp.abs(fieldline_SIMSOPT_t[1] - fieldline_ESSOS_t[1])/(jnp.abs(fieldline_SIMSOPT_t[1])+1e-12) - relative_error_z = jnp.abs(fieldline_SIMSOPT_t[2] - fieldline_ESSOS_t[2])/(jnp.abs(fieldline_SIMSOPT_t[2])+1e-12) - average_relative_error.append((relative_error_x + relative_error_y + relative_error_z)/3) - average_relative_error = jnp.array(average_relative_error) - relative_error_fieldlines_SIMSOPT_vs_ESSOS.append(average_relative_error) - plt.plot(jnp.linspace(0, tmax_fl, len(average_relative_error))[1:], average_relative_error[1:], label=f'Fieldline {j}') - plt.legend() - plt.xlabel('Time') - plt.ylabel('Relative Error') - plt.yscale('log') - plt.tight_layout() - plt.savefig(os.path.join(output_dir, f'relative_error_fieldlines_SIMSOPT_vs_ESSOS_tolerance{trace_tolerance_SIMSOPT_array[i]}.pdf'), dpi=150) - plt.close() - - # relative_error_fieldlines_SIMSOPT_vs_ESSOS = jnp.array(relative_error_fieldlines_SIMSOPT_vs_ESSOS) - # print(f"Relative difference between SIMSOPT and ESSOS fieldlines={relative_error_fieldlines_SIMSOPT_vs_ESSOS}") - relative_error_array.append(relative_error_fieldlines_SIMSOPT_vs_ESSOS) - - plt.figure() - for j in range(nfieldlines): - R_SIMSOPT = jnp.sqrt(fieldlines_SIMSOPT[j][:,1]**2+fieldlines_SIMSOPT[j][:,2]**2) - phi_SIMSOPT = jnp.arctan2(fieldlines_SIMSOPT[j][:,2], fieldlines_SIMSOPT[j][:,1]) - Z_SIMSOPT = fieldlines_SIMSOPT[j][:,3] - - R_ESSOS = jnp.sqrt(fieldlines_ESSOS_interp[j][:,0]**2+fieldlines_ESSOS_interp[j][:,1]**2) - phi_ESSOS = jnp.arctan2(fieldlines_ESSOS_interp[j][:,1], fieldlines_ESSOS_interp[j][:,0]) - Z_ESSOS = fieldlines_ESSOS_interp[j][:,2] - - plt.plot(R_SIMSOPT, Z_SIMSOPT, '-', linewidth=2.5, label=f'SIMSOPT {j}') - plt.plot(R_ESSOS, Z_ESSOS, '--', linewidth=2.5, label=f'ESSOS {j}') - plt.legend() - plt.xlabel('R') - plt.ylabel('Z') - plt.savefig(os.path.join(output_dir,f'fieldlines_SIMSOPT_vs_ESSOS_tolerance{trace_tolerance_SIMSOPT_array[i]}.pdf'), dpi=150) - plt.close() - -# Calculate RMS error for each tolerance -rms_error_array = jnp.array([[jnp.sqrt(jnp.mean(jnp.square(jnp.array(error)))) for error in relative_error] for relative_error in relative_error_array]) - -# Plot RMS error in a bar chart -plt.figure() -bar_width = 0.15 -x = jnp.arange(len(trace_tolerance_SIMSOPT_array)) -for i in range(rms_error_array.shape[1]): - plt.bar(x + i * bar_width, rms_error_array[:, i], bar_width, label=f'Fieldline {i}') -plt.xlabel('Tracing Tolerance of SIMSOPT') -plt.ylabel('RMS Error') -plt.yscale('log') -plt.xticks(x + bar_width * (rms_error_array.shape[1] - 1) / 2, [f'Tol={tol}' for tol in trace_tolerance_SIMSOPT_array], rotation=45) -plt.legend() -plt.tight_layout() -plt.savefig(os.path.join(output_dir, 'rms_error_fieldlines_SIMSOPT_vs_ESSOS.pdf'), dpi=150) -plt.close() - -# Calculate maximum error for each tolerance -max_error_array = jnp.array([[jnp.max(jnp.array(error)) for error in relative_error] for relative_error in relative_error_array]) -# Plot maximum error in a bar chart -plt.figure() -bar_width = 0.15 -x = jnp.arange(len(trace_tolerance_SIMSOPT_array)) -for i in range(max_error_array.shape[1]): - plt.bar(x + i * bar_width, max_error_array[:, i], bar_width, label=f'Fieldline {i}') -plt.xlabel('Tracing Tolerance of SIMSOPT') -plt.ylabel('Maximum Error') -plt.yscale('log') -plt.xticks(x + bar_width * (max_error_array.shape[1] - 1) / 2, [f'Tol={tol}' for tol in trace_tolerance_SIMSOPT_array], rotation=45) -plt.legend() -plt.tight_layout() -plt.savefig(os.path.join(output_dir, 'max_error_fieldlines_SIMSOPT_vs_ESSOS.pdf'), dpi=150) -plt.close() - -# Calculate mean error for each tolerance -mean_error_array = jnp.array([[jnp.mean(jnp.array(error)) for error in relative_error] for relative_error in relative_error_array]) -# Plot mean error in a bar chart -plt.figure() -bar_width = 0.15 -x = jnp.arange(len(trace_tolerance_SIMSOPT_array)) -for i in range(mean_error_array.shape[1]): - plt.bar(x + i * bar_width, mean_error_array[:, i], bar_width, label=f'Fieldline {i}') -plt.xlabel('Tracing Tolerance of SIMSOPT') -plt.ylabel('Mean Error') -plt.yscale('log') -plt.xticks(x + bar_width * (mean_error_array.shape[1] - 1) / 2, [f'Tol={tol}' for tol in trace_tolerance_SIMSOPT_array], rotation=45) -plt.legend() -plt.tight_layout() -plt.savefig(os.path.join(output_dir, 'mean_error_fieldlines_SIMSOPT_vs_ESSOS.pdf'), dpi=150) -plt.close() diff --git a/examples/comparisons_SIMSOPT/fullorbit_SIMSOPT_vs_ESSOS.py b/examples/comparisons_SIMSOPT/fullorbit_SIMSOPT_vs_ESSOS.py deleted file mode 100644 index c1b2aba1..00000000 --- a/examples/comparisons_SIMSOPT/fullorbit_SIMSOPT_vs_ESSOS.py +++ /dev/null @@ -1,259 +0,0 @@ -import os -import time -import jax.numpy as jnp -from jax import block_until_ready, random -from simsopt import load -from simsopt.field import (particles_to_vtk, trace_particles, plot_poincare_data) -from essos.coils import Coils_from_simsopt -from essos.constants import PROTON_MASS, ONE_EV -from essos.dynamics import Tracing, Particles -from essos.fields import BiotSavart as BiotSavart_essos -import matplotlib.pyplot as plt - -tmax_full = 1e-5 -nparticles = 3 -axis_shft=0.02 -R0 = jnp.linspace(1.2125346+axis_shft, 1.295-axis_shft, nparticles) -trace_tolerance_SIMSOPT_array = [1e-3, 1e-5, 1e-7, 1e-9]#, 1e-11] -trace_tolerance_ESSOS = 1e-5 -mass=PROTON_MASS -energy=5000*ONE_EV -model_ESSOS_array = ['FullOrbit', 'FullOrbit_Boris'] - -output_dir = os.path.join(os.path.dirname(__file__), 'output') -if not os.path.exists(output_dir): - os.makedirs(output_dir) - -nfp=2 -LandremanPaulQA_json_file = os.path.join(os.path.dirname(__file__), '..', 'input_files', 'SIMSOPT_biot_savart_LandremanPaulQA.json') -field_simsopt = load(LandremanPaulQA_json_file) -field_essos = BiotSavart_essos(Coils_from_simsopt(LandremanPaulQA_json_file, nfp)) - -Z0 = jnp.zeros(nparticles) -phi0 = jnp.zeros(nparticles) -initial_xyz=jnp.array([R0*jnp.cos(phi0), R0*jnp.sin(phi0), Z0]).T -initial_vparallel_over_v = random.uniform(random.PRNGKey(42), (nparticles,), minval=-1, maxval=1) - - -phis_poincare = [(i/4)*(2*jnp.pi/nfp) for i in range(4)] - -particles = Particles(initial_xyz=initial_xyz, initial_vparallel_over_v=initial_vparallel_over_v, mass=mass, energy=energy, field=field_essos) - -# Trace in SIMSOPT -time_SIMSOPT_array = [] -trajectories_SIMSOPT_array = [] -avg_steps_SIMSOPT = 0 -relative_energy_error_SIMSOPT_array = [] -print(f'Output being saved to {output_dir}') -print(f'SIMSOPT LandremanPaulQA json file location: {LandremanPaulQA_json_file}') -for trace_tolerance_SIMSOPT in trace_tolerance_SIMSOPT_array: - print(f' Tracing SIMSOPT full orbit with tolerance={trace_tolerance_SIMSOPT}') - t1 = time.time() - trajectories_SIMSOPT_this_tolerance, trajectories_SIMSOPT_phi_hits = block_until_ready(trace_particles( - field=field_simsopt, xyz_inits=particles.initial_xyz, mass=particles.mass, - parallel_speeds=particles.initial_vparallel, tmax=tmax_full, mode='full', - charge=particles.charge, Ekin=particles.energy, tol=trace_tolerance_SIMSOPT)) - time_SIMSOPT_array.append(time.time()-t1) - avg_steps_SIMSOPT += sum([len(l) for l in trajectories_SIMSOPT_this_tolerance])//nparticles - print(f" Time for SIMSOPT tracing={time.time()-t1:.3f}s. Avg num steps={avg_steps_SIMSOPT}") - trajectories_SIMSOPT_array.append(trajectories_SIMSOPT_this_tolerance) - - relative_energy_error_SIMSOPT_array.append([jnp.abs(mass*(trajectory[:,4]**2+trajectory[:,5]**2+trajectory[:,6]**2)/2-particles.energy)/particles.energy - for trajectory in trajectories_SIMSOPT_this_tolerance]) - -particles_to_vtk(trajectories_SIMSOPT_this_tolerance, os.path.join(output_dir,f'full_orbit_SIMSOPT')) - - -# Trace in ESSOS -num_steps_essos = int(jnp.max(jnp.array([len(trajectories_SIMSOPT[0]) for trajectories_SIMSOPT in trajectories_SIMSOPT_array]))) -time_essos = jnp.linspace(0, tmax_full, num_steps_essos) - - -tracing_array = [] -trajectories_ESSOS_array = [] -time_ESSOS_array = [] -for model_ESSOS in model_ESSOS_array: - print(f'Tracing ESSOS full orbit '+('Boris' if model_ESSOS=='FullOrbit_Boris' else f'with tolerance={trace_tolerance_ESSOS}')+f' and plotting the result.') - t1 = time.time() - tracing = block_until_ready(Tracing(field=field_essos, model=model_ESSOS, particles=particles, - maxtime=tmax_full, timesteps=num_steps_essos, tol_step_size=trace_tolerance_ESSOS)) - trajectories_ESSOS = tracing.trajectories - time_ESSOS = time.time()-t1 - print(f" Time for ESSOS tracing={time.time()-t1:.3f}s "+('Boris' if model_ESSOS=='FullOrbit_Boris' else f'')+f". Num steps={len(trajectories_ESSOS[0])}") - tracing.to_vtk(os.path.join(output_dir,f'full_orbit'+('_boris' if model_ESSOS=='FullOrbit_Boris' else '')+'_ESSOS')) - tracing_array.append(tracing) - trajectories_ESSOS_array.append(trajectories_ESSOS) - time_ESSOS_array.append(time_ESSOS) - -print('Plotting the results to output directory...') -plt.figure() -SIMSOPT_energy_interp_this_particle = jnp.zeros((len(trace_tolerance_SIMSOPT_array), nparticles, len(trajectories_SIMSOPT_array[-1][-1][:,0]))) -for j in range(nparticles): - for i, relative_energy_error_SIMSOPT in enumerate(relative_energy_error_SIMSOPT_array): - SIMSOPT_energy_interp_this_particle = SIMSOPT_energy_interp_this_particle.at[i,j].set(jnp.interp(trajectories_SIMSOPT_array[-1][-1][:,0], trajectories_SIMSOPT_array[i][j][:,0], relative_energy_error_SIMSOPT[j][:])) -for i, SIMSOPT_energy_interp in enumerate(SIMSOPT_energy_interp_this_particle): - plt.plot(trajectories_SIMSOPT_array[-1][-1][4:,0], jnp.mean(SIMSOPT_energy_interp, axis=0)[4:], '--', label=f'SIMSOPT Tol={trace_tolerance_SIMSOPT_array[i]}') -for model_ESSOS, tracing, trajectories_ESSOS in zip(model_ESSOS_array, tracing_array, trajectories_ESSOS_array): - relative_energy_error_ESSOS = jnp.abs(tracing.energy-particles.energy)/particles.energy - plt.plot(time_essos[2:], jnp.mean(relative_energy_error_ESSOS, axis=0)[2:], '-', label=f'ESSOS'+(' Boris' if model_ESSOS=='FullOrbit_Boris' else f' Tol={trace_tolerance_ESSOS}')) -plt.legend() -plt.yscale('log') -plt.xlabel('Time (s)') -plt.ylabel('Average Relative Energy Error') -plt.tight_layout() -plt.savefig(os.path.join(output_dir, f'relative_energy_error_full_orbit_SIMSOPT_vs_ESSOS.pdf'), dpi=150) -plt.close() - -labels = [f'SIMSOPT Tol={tol}' for tol in trace_tolerance_SIMSOPT_array] -times = time_SIMSOPT_array -plt.figure() -for model_ESSOS, tracing, trajectories_ESSOS, time_ESSOS in zip(model_ESSOS_array, tracing_array, trajectories_ESSOS_array, time_ESSOS_array): - # Plot time comparison in a bar chart - labels += ([f'ESSOS Boris Algorithm'] if model_ESSOS=='FullOrbit_Boris' else [f'ESSOS Tol={trace_tolerance_ESSOS}']) - times += [time_ESSOS] -bars = plt.bar(labels, times, color=['blue']*len(trace_tolerance_SIMSOPT_array) + ['red', 'orange'], edgecolor=['black']*len(trace_tolerance_SIMSOPT_array) + ['black']*2, hatch=['//']*len(trace_tolerance_SIMSOPT_array) + ['|']*2) -plt.xlabel('Tracing Tolerance of SIMSOPT') -plt.ylabel('Time (s)') -plt.xticks(rotation=45) -plt.tight_layout() -blue_patch = plt.Line2D([0], [0], color='blue', lw=4, label='SIMSOPT', linestyle='--') -red_patch = plt.Line2D([0], [0], color='red', lw=4, label=f'ESSOS', linestyle='-') -orange_patch = plt.Line2D([0], [0], color='orange', lw=4, label=f'ESSOS\nBoris Algorithm') -plt.legend(handles=[blue_patch, red_patch, orange_patch]) -plt.savefig(os.path.join(output_dir, 'times_full_orbit'+('_boris' if model_ESSOS=='FullOrbit_Boris' else '')+'_SIMSOPT_vs_ESSOS.pdf'), dpi=150) -plt.close() - -def interpolate_ESSOS_to_SIMSOPT(trajectory_SIMSOPT, trajectory_ESSOS): - time_SIMSOPT = jnp.array(trajectory_SIMSOPT)[:, 0] # Time values from full orbit SIMSOPT - # coords_SIMSOPT = jnp.array(trajectory_SIMSOPT)[:, 1:] # Coordinates (x, y, z) from full orbit SIMSOPT - coords_ESSOS = jnp.array(trajectory_ESSOS) - interp_x = jnp.interp(time_SIMSOPT, time_essos, coords_ESSOS[:, 0]) - interp_y = jnp.interp(time_SIMSOPT, time_essos, coords_ESSOS[:, 1]) - interp_z = jnp.interp(time_SIMSOPT, time_essos, coords_ESSOS[:, 2]) - interp_vx = jnp.interp(time_SIMSOPT, time_essos, coords_ESSOS[:, 3]) - interp_vy = jnp.interp(time_SIMSOPT, time_essos, coords_ESSOS[:, 4]) - interp_vz = jnp.interp(time_SIMSOPT, time_essos, coords_ESSOS[:, 5]) - coords_ESSOS_interp = jnp.column_stack([ interp_x, interp_y, interp_z, interp_vx, interp_vy, interp_vz]) - return coords_ESSOS_interp - -for model_ESSOS, tracing, trajectories_ESSOS, time_ESSOS in zip(model_ESSOS_array, tracing_array, trajectories_ESSOS_array, time_ESSOS_array): - - relative_error_array = [] - for i, trajectories_SIMSOPT in enumerate(trajectories_SIMSOPT_array): - trajectories_ESSOS_interp = [interpolate_ESSOS_to_SIMSOPT(trajectories_SIMSOPT[i], trajectories_ESSOS[i]) for i in range(nparticles)] - tracing.trajectories = trajectories_ESSOS_interp - if i==len(trace_tolerance_SIMSOPT_array)-1: tracing.to_vtk(os.path.join(output_dir,f'full_orbit'+('_boris' if model_ESSOS=='FullOrbit_Boris' else '')+'_ESSOS_interp')) - - relative_error_trajectories_SIMSOPT_vs_ESSOS = [] - plt.figure() - for j in range(nparticles): - this_trajectory_SIMSOPT = jnp.array(trajectories_SIMSOPT[j])[:,1:] - this_trajectory_ESSOS = trajectories_ESSOS_interp[j] - average_relative_error = [] - for trajectory_SIMSOPT_t, trajectory_ESSOS_t in zip(this_trajectory_SIMSOPT, this_trajectory_ESSOS): - relative_error_x = jnp.abs(trajectory_SIMSOPT_t[0] - trajectory_ESSOS_t[0])/(jnp.abs(trajectory_SIMSOPT_t[0])+1e-12) - relative_error_y = jnp.abs(trajectory_SIMSOPT_t[1] - trajectory_ESSOS_t[1])/(jnp.abs(trajectory_SIMSOPT_t[1])+1e-12) - relative_error_z = jnp.abs(trajectory_SIMSOPT_t[2] - trajectory_ESSOS_t[2])/(jnp.abs(trajectory_SIMSOPT_t[2])+1e-12) - relative_error_vx = jnp.abs(trajectory_SIMSOPT_t[3] - trajectory_ESSOS_t[3])/(jnp.abs(trajectory_SIMSOPT_t[3])+1e-12) - relative_error_vy = jnp.abs(trajectory_SIMSOPT_t[3] - trajectory_ESSOS_t[3])/(jnp.abs(trajectory_SIMSOPT_t[4])+1e-12) - relative_error_vz = jnp.abs(trajectory_SIMSOPT_t[3] - trajectory_ESSOS_t[3])/(jnp.abs(trajectory_SIMSOPT_t[5])+1e-12) - average_relative_error.append((relative_error_x + relative_error_y + relative_error_z + relative_error_vx + relative_error_vy + relative_error_vz)/6) - average_relative_error = jnp.array(average_relative_error) - relative_error_trajectories_SIMSOPT_vs_ESSOS.append(average_relative_error) - plt.plot(jnp.linspace(0, tmax_full, len(average_relative_error))[1:], average_relative_error[1:], label=f'Particle {1+j}') - plt.legend() - plt.xlabel('Time') - plt.ylabel('Relative Error') - plt.yscale('log') - plt.tight_layout() - plt.savefig(os.path.join(output_dir, f'relative_error_full_orbit'+('_boris' if model_ESSOS=='FullOrbit_Boris' else '')+f'_SIMSOPT_vs_ESSOS_tolerance{trace_tolerance_SIMSOPT_array[i]}.pdf'), dpi=150) - plt.close() - - relative_error_array.append(relative_error_trajectories_SIMSOPT_vs_ESSOS) - - plt.figure() - for j in range(nparticles): - R_SIMSOPT = jnp.sqrt(trajectories_SIMSOPT[j][:,1]**2+trajectories_SIMSOPT[j][:,2]**2) - phi_SIMSOPT = jnp.arctan2(trajectories_SIMSOPT[j][:,2], trajectories_SIMSOPT[j][:,1]) - Z_SIMSOPT = trajectories_SIMSOPT[j][:,3] - - R_ESSOS = jnp.sqrt(trajectories_ESSOS_interp[j][:,0]**2+trajectories_ESSOS_interp[j][:,1]**2) - phi_ESSOS = jnp.arctan2(trajectories_ESSOS_interp[j][:,1], trajectories_ESSOS_interp[j][:,0]) - Z_ESSOS = trajectories_ESSOS_interp[j][:,2] - - plt.plot(R_SIMSOPT, Z_SIMSOPT, '-', linewidth=2.5, label=f'SIMSOPT {1+j}') - plt.plot(R_ESSOS, Z_ESSOS, '--', linewidth=2.5, label=f'ESSOS {1+j}') - plt.legend() - plt.xlabel('R') - plt.ylabel('Z') - plt.tight_layout() - plt.savefig(os.path.join(output_dir,f'full_orbit'+('_boris' if model_ESSOS=='FullOrbit_Boris' else '')+f'_RZ_SIMSOPT_vs_ESSOS_tolerance{trace_tolerance_SIMSOPT_array[i]}.pdf'), dpi=150) - plt.close() - - plt.figure() - for j in range(nparticles): - time_SIMSOPT = jnp.array(trajectories_SIMSOPT[j][:,0]) - vx_SIMSOPT = jnp.array(trajectories_SIMSOPT[j][:,4]) - vx_ESSOS = jnp.array(trajectories_ESSOS_interp[j][:,3]) - # plt.plot(time_SIMSOPT, jnp.abs((vx_SIMSOPT-vx_ESSOS)/vx_SIMSOPT), '-', linewidth=2.5, label=f'Particle {1+j}') - plt.plot(time_SIMSOPT, vx_SIMSOPT/particles.total_speed, '-', linewidth=2.5, label=f'SIMSOPT {1+j}') - plt.plot(time_SIMSOPT, vx_ESSOS/particles.total_speed, '--', linewidth=2.5, label=f'ESSOS {1+j}') - plt.legend() - plt.xlabel('Time (s)') - plt.ylabel(r'$v_x/v$') - # plt.yscale('log') - plt.tight_layout() - plt.savefig(os.path.join(output_dir,f'full_orbit'+('_boris' if model_ESSOS=='FullOrbit_Boris' else '')+f'_vx_SIMSOPT_vs_ESSOS_tolerance{trace_tolerance_SIMSOPT_array[i]}.pdf'), dpi=150) - plt.close() - - # Calculate RMS error for each tolerance - rms_error_array = jnp.array([[jnp.sqrt(jnp.mean(jnp.square(jnp.array(error)))) for error in relative_error] for relative_error in relative_error_array]) - - # Plot RMS error in a bar chart - plt.figure() - bar_width = 0.15 - x = jnp.arange(len(trace_tolerance_SIMSOPT_array)) - for i in range(rms_error_array.shape[1]): - plt.bar(x + i * bar_width, rms_error_array[:, i], bar_width, label=f'Particle {1+i}') - plt.xlabel('Tracing Tolerance of SIMSOPT') - plt.ylabel('RMS Error') - plt.yscale('log') - plt.xticks(x + bar_width * (rms_error_array.shape[1] - 1) / 2, [f'Tol={tol}' for tol in trace_tolerance_SIMSOPT_array], rotation=45) - plt.legend() - plt.tight_layout() - plt.savefig(os.path.join(output_dir, 'rms_error_full_orbit'+('_boris' if model_ESSOS=='FullOrbit_Boris' else '')+'_SIMSOPT_vs_ESSOS.pdf'), dpi=150) - plt.close() - - # Calculate maximum error for each tolerance - max_error_array = jnp.array([[jnp.max(jnp.array(error)) for error in relative_error] for relative_error in relative_error_array]) - # Plot maximum error in a bar chart - plt.figure() - bar_width = 0.15 - x = jnp.arange(len(trace_tolerance_SIMSOPT_array)) - for i in range(max_error_array.shape[1]): - plt.bar(x + i * bar_width, max_error_array[:, i], bar_width, label=f'Particle {1+i}') - plt.xlabel('Tracing Tolerance of SIMSOPT') - plt.ylabel('Maximum Error') - plt.yscale('log') - plt.xticks(x + bar_width * (max_error_array.shape[1] - 1) / 2, [f'Tol={tol}' for tol in trace_tolerance_SIMSOPT_array], rotation=45) - plt.legend() - plt.tight_layout() - plt.savefig(os.path.join(output_dir, 'max_error_full_orbit'+('_boris' if model_ESSOS=='FullOrbit_Boris' else '')+'_SIMSOPT_vs_ESSOS.pdf'), dpi=150) - plt.close() - - # Calculate mean error for each tolerance - mean_error_array = jnp.array([[jnp.mean(jnp.array(error)) for error in relative_error] for relative_error in relative_error_array]) - # Plot mean error in a bar chart - plt.figure() - bar_width = 0.15 - x = jnp.arange(len(trace_tolerance_SIMSOPT_array)) - for i in range(mean_error_array.shape[1]): - plt.bar(x + i * bar_width, mean_error_array[:, i], bar_width, label=f'Particle {1+i}') - plt.xlabel('Tracing Tolerance of SIMSOPT') - plt.ylabel('Mean Error') - plt.yscale('log') - plt.xticks(x + bar_width * (mean_error_array.shape[1] - 1) / 2, [f'Tol={tol}' for tol in trace_tolerance_SIMSOPT_array], rotation=45) - plt.legend() - plt.tight_layout() - plt.savefig(os.path.join(output_dir, 'mean_error_full_orbit'+('_boris' if model_ESSOS=='FullOrbit_Boris' else '')+'_SIMSOPT_vs_ESSOS.pdf'), dpi=150) - plt.close() diff --git a/examples/comparisons_SIMSOPT/guiding_center_SIMSOPT_vs_ESSOS.py b/examples/comparisons_SIMSOPT/guiding_center_SIMSOPT_vs_ESSOS.py deleted file mode 100644 index eb102a77..00000000 --- a/examples/comparisons_SIMSOPT/guiding_center_SIMSOPT_vs_ESSOS.py +++ /dev/null @@ -1,248 +0,0 @@ -import os -import time -import jax.numpy as jnp -from jax import block_until_ready, random -from simsopt import load -from simsopt.field import (particles_to_vtk, trace_particles, plot_poincare_data) -from essos.coils import Coils_from_simsopt -from essos.constants import PROTON_MASS, ONE_EV -from essos.dynamics import Tracing, Particles -from essos.fields import BiotSavart as BiotSavart_essos -import matplotlib.pyplot as plt - -tmax_gc = 1e-4 -nparticles = 5 -axis_shft=0.02 -R0 = jnp.linspace(1.2125346+axis_shft, 1.295-axis_shft, nparticles) -trace_tolerance_SIMSOPT_array = [1e-5, 1e-7, 1e-9, 1e-11] -trace_tolerance_ESSOS = 1e-7 -mass=PROTON_MASS -energy=5000*ONE_EV - -output_dir = os.path.join(os.path.dirname(__file__), 'output') -if not os.path.exists(output_dir): - os.makedirs(output_dir) - -nfp=2 -LandremanPaulQA_json_file = os.path.join(os.path.dirname(__file__), '..', 'input_files', 'SIMSOPT_biot_savart_LandremanPaulQA.json') -field_simsopt = load(LandremanPaulQA_json_file) -field_essos = BiotSavart_essos(Coils_from_simsopt(LandremanPaulQA_json_file, nfp)) - -Z0 = jnp.zeros(nparticles) -phi0 = jnp.zeros(nparticles) -initial_xyz=jnp.array([R0*jnp.cos(phi0), R0*jnp.sin(phi0), Z0]).T -initial_vparallel_over_v = random.uniform(random.PRNGKey(42), (nparticles,), minval=-1, maxval=1) - -phis_poincare = [(i/4)*(2*jnp.pi/nfp) for i in range(4)] - -particles = Particles(initial_xyz=initial_xyz, initial_vparallel_over_v=initial_vparallel_over_v, mass=mass, energy=energy) - -# Trace in SIMSOPT -time_SIMSOPT_array = [] -trajectories_SIMSOPT_array = [] -avg_steps_SIMSOPT = 0 -relative_energy_error_SIMSOPT_array = [] -print(f'Output being saved to {output_dir}') -print(f'SIMSOPT LandremanPaulQA json file location: {LandremanPaulQA_json_file}') -for trace_tolerance_SIMSOPT in trace_tolerance_SIMSOPT_array: - print(f'Tracing SIMSOPT guiding center with tolerance={trace_tolerance_SIMSOPT}') - t1 = time.time() - trajectories_SIMSOPT_this_tolerance, trajectories_SIMSOPT_phi_hits = block_until_ready(trace_particles( - field=field_simsopt, xyz_inits=particles.initial_xyz, mass=particles.mass, - parallel_speeds=particles.initial_vparallel, tmax=tmax_gc, mode='gc_vac', - charge=particles.charge, Ekin=particles.energy, tol=trace_tolerance_SIMSOPT)) - time_SIMSOPT_array.append(time.time()-t1) - avg_steps_SIMSOPT += sum([len(l) for l in trajectories_SIMSOPT_this_tolerance])//nparticles - print(f" Time for SIMSOPT tracing={time.time()-t1:.3f}s. Avg num steps={avg_steps_SIMSOPT}") - trajectories_SIMSOPT_array.append(trajectories_SIMSOPT_this_tolerance) - - relative_energy_SIMSOPT = [] - for i, trajectory in enumerate(trajectories_SIMSOPT_this_tolerance): - xyz = jnp.asarray(trajectory[:, 1:4]) - vpar = trajectory[:, 4] - field_simsopt.set_points(xyz) - AbsB = field_simsopt.AbsB()[:,0] - mu = (particles.energy - particles.mass*vpar[0]**2/2)/AbsB[0] - relative_energy_SIMSOPT.append(jnp.abs(particles.mass*vpar**2/2+mu*AbsB-particles.energy)/particles.energy) - relative_energy_error_SIMSOPT_array.append(relative_energy_SIMSOPT) - -particles_to_vtk(trajectories_SIMSOPT_this_tolerance, os.path.join(output_dir,f'guiding_center_SIMSOPT')) - -# Trace in ESSOS -num_steps_essos = int(jnp.mean(jnp.array([len(trajectories_SIMSOPT[0]) for trajectories_SIMSOPT in trajectories_SIMSOPT_array]))) -time_essos = jnp.linspace(0, tmax_gc, num_steps_essos) - -print(f'Tracing ESSOS guiding center with tolerance={trace_tolerance_ESSOS}') -t1 = time.time() -tracing = block_until_ready(Tracing(field=field_essos, model='GuidingCenter', particles=particles, - maxtime=tmax_gc, timesteps=num_steps_essos, tol_step_size=trace_tolerance_ESSOS)) -trajectories_ESSOS = tracing.trajectories -time_ESSOS = time.time()-t1 -print(f" Time for ESSOS tracing={time.time()-t1:.3f}s. Num steps={len(trajectories_ESSOS[0])}") -tracing.to_vtk(os.path.join(output_dir,f'guiding_center_ESSOS')) - -relative_energy_error_ESSOS = jnp.abs(tracing.energy-particles.energy)/particles.energy - -print('Plotting the results to output directory...') -plt.figure() -SIMSOPT_energy_interp_this_particle = jnp.zeros((len(trace_tolerance_SIMSOPT_array), nparticles, len(trajectories_SIMSOPT_array[-1][-1][:,0]))) -for j in range(nparticles): - for i, relative_energy_error_SIMSOPT in enumerate(relative_energy_error_SIMSOPT_array): - SIMSOPT_energy_interp_this_particle = SIMSOPT_energy_interp_this_particle.at[i,j].set(jnp.interp(trajectories_SIMSOPT_array[-1][-1][:,0], trajectories_SIMSOPT_array[i][j][:,0], relative_energy_error_SIMSOPT[j][:])) -plt.plot(time_essos[2:], jnp.mean(relative_energy_error_ESSOS, axis=0)[2:], '-', label=f'ESSOS Tol={trace_tolerance_ESSOS}') -for i, SIMSOPT_energy_interp in enumerate(SIMSOPT_energy_interp_this_particle): - plt.plot(trajectories_SIMSOPT_array[-1][-1][4:,0], jnp.mean(SIMSOPT_energy_interp, axis=0)[4:], '--', label=f'SIMSOPT Tol={trace_tolerance_SIMSOPT_array[i]}') -plt.legend() -plt.yscale('log') -plt.xlabel('Time (s)') -plt.ylabel('Average Relative Energy Error') -plt.tight_layout() -plt.savefig(os.path.join(output_dir, f'relative_energy_error_guiding_center_SIMSOPT_vs_ESSOS.pdf'), dpi=150) -plt.close() - -# Plot time comparison in a bar chart -labels = [f'SIMSOPT\nTol={tol}' for tol in trace_tolerance_SIMSOPT_array] + [f'ESSOS\nTol={trace_tolerance_ESSOS}'] -times = time_SIMSOPT_array + [time_ESSOS] -plt.figure() -bars = plt.bar(labels, times, color=['blue']*len(trace_tolerance_SIMSOPT_array) + ['red'], edgecolor=['black']*len(trace_tolerance_SIMSOPT_array) + ['black'], hatch=['//']*len(trace_tolerance_SIMSOPT_array) + ['|']) -plt.xlabel('Tracing Tolerance of SIMSOPT') -plt.ylabel('Time (s)') -plt.xticks(rotation=45) -plt.tight_layout() -blue_patch = plt.Line2D([0], [0], color='blue', lw=4, label='SIMSOPT', linestyle='--') -orange_patch = plt.Line2D([0], [0], color='red', lw=4, label=f'ESSOS', linestyle='-') -plt.legend(handles=[blue_patch, orange_patch]) -plt.savefig(os.path.join(output_dir, 'times_guiding_center_SIMSOPT_vs_ESSOS.pdf'), dpi=150) -plt.close() - -def interpolate_ESSOS_to_SIMSOPT(trajectory_SIMSOPT, trajectory_ESSOS): - time_SIMSOPT = jnp.array(trajectory_SIMSOPT)[:, 0] # Time values from guiding center SIMSOPT - # coords_SIMSOPT = jnp.array(trajectory_SIMSOPT)[:, 1:] # Coordinates (x, y, z) from guiding center SIMSOPT - coords_ESSOS = jnp.array(trajectory_ESSOS) - - interp_x = jnp.interp(time_SIMSOPT, time_essos, coords_ESSOS[:, 0]) - interp_y = jnp.interp(time_SIMSOPT, time_essos, coords_ESSOS[:, 1]) - interp_z = jnp.interp(time_SIMSOPT, time_essos, coords_ESSOS[:, 2]) - interp_v = jnp.interp(time_SIMSOPT, time_essos, coords_ESSOS[:, 3]) - - coords_ESSOS_interp = jnp.column_stack([ interp_x, interp_y, interp_z, interp_v]) - - return coords_ESSOS_interp - -relative_error_array = [] -for i, trajectories_SIMSOPT in enumerate(trajectories_SIMSOPT_array): - trajectories_ESSOS_interp = [interpolate_ESSOS_to_SIMSOPT(trajectories_SIMSOPT[i], trajectories_ESSOS[i]) for i in range(nparticles)] - tracing.trajectories = trajectories_ESSOS_interp - if i==len(trace_tolerance_SIMSOPT_array)-1: tracing.to_vtk(os.path.join(output_dir,f'guiding_center_ESSOS_interp')) - - relative_error_trajectories_SIMSOPT_vs_ESSOS = [] - plt.figure() - for j in range(nparticles): - this_trajectory_SIMSOPT = jnp.array(trajectories_SIMSOPT[j])[:,1:] - this_trajectory_ESSOS = trajectories_ESSOS_interp[j] - average_relative_error = [] - for trajectory_SIMSOPT_t, trajectory_ESSOS_t in zip(this_trajectory_SIMSOPT, this_trajectory_ESSOS): - relative_error_x = jnp.abs(trajectory_SIMSOPT_t[0] - trajectory_ESSOS_t[0])/(jnp.abs(trajectory_SIMSOPT_t[0])+1e-12) - relative_error_y = jnp.abs(trajectory_SIMSOPT_t[1] - trajectory_ESSOS_t[1])/(jnp.abs(trajectory_SIMSOPT_t[1])+1e-12) - relative_error_z = jnp.abs(trajectory_SIMSOPT_t[2] - trajectory_ESSOS_t[2])/(jnp.abs(trajectory_SIMSOPT_t[2])+1e-12) - relative_error_v = jnp.abs(trajectory_SIMSOPT_t[3] - trajectory_ESSOS_t[3])/(jnp.abs(trajectory_SIMSOPT_t[3])+1e-12) - average_relative_error.append((relative_error_x + relative_error_y + relative_error_z + relative_error_v)/4) - average_relative_error = jnp.array(average_relative_error) - relative_error_trajectories_SIMSOPT_vs_ESSOS.append(average_relative_error) - plt.plot(jnp.linspace(0, tmax_gc, len(average_relative_error))[1:], average_relative_error[1:], label=f'Particle {1+j}') - plt.legend() - plt.xlabel('Time') - plt.ylabel('Relative Error') - plt.yscale('log') - plt.tight_layout() - plt.savefig(os.path.join(output_dir, f'relative_error_guiding_center_SIMSOPT_vs_ESSOS_tolerance{trace_tolerance_SIMSOPT_array[i]}.pdf'), dpi=150) - plt.close() - - relative_error_array.append(relative_error_trajectories_SIMSOPT_vs_ESSOS) - - plt.figure() - for j in range(nparticles): - R_SIMSOPT = jnp.sqrt(trajectories_SIMSOPT[j][:,1]**2+trajectories_SIMSOPT[j][:,2]**2) - phi_SIMSOPT = jnp.arctan2(trajectories_SIMSOPT[j][:,2], trajectories_SIMSOPT[j][:,1]) - Z_SIMSOPT = trajectories_SIMSOPT[j][:,3] - - R_ESSOS = jnp.sqrt(trajectories_ESSOS_interp[j][:,0]**2+trajectories_ESSOS_interp[j][:,1]**2) - phi_ESSOS = jnp.arctan2(trajectories_ESSOS_interp[j][:,1], trajectories_ESSOS_interp[j][:,0]) - Z_ESSOS = trajectories_ESSOS_interp[j][:,2] - - plt.plot(R_SIMSOPT, Z_SIMSOPT, '-', linewidth=2.5, label=f'SIMSOPT {1+j}') - plt.plot(R_ESSOS, Z_ESSOS, '--', linewidth=2.5, label=f'ESSOS {1+j}') - plt.legend() - plt.xlabel('R') - plt.ylabel('Z') - plt.tight_layout() - plt.savefig(os.path.join(output_dir,f'guiding_center_RZ_SIMSOPT_vs_ESSOS_tolerance{trace_tolerance_SIMSOPT_array[i]}.pdf'), dpi=150) - plt.close() - - plt.figure() - for j in range(nparticles): - time_SIMSOPT = jnp.array(trajectories_SIMSOPT[j][:,0]) - vpar_SIMSOPT = jnp.array(trajectories_SIMSOPT[j][:,4]) - vpar_ESSOS = jnp.array(trajectories_ESSOS_interp[j][:,3]) - # plt.plot(time_SIMSOPT, jnp.abs((vpar_SIMSOPT-vpar_ESSOS)/vpar_SIMSOPT), '-', linewidth=2.5, label=f'Particle {1+j}') - plt.plot(time_SIMSOPT, vpar_SIMSOPT, '-', linewidth=2.5, label=f'SIMSOPT {1+j}') - plt.plot(time_SIMSOPT, vpar_ESSOS, '--', linewidth=2.5, label=f'ESSOS {1+j}') - plt.legend() - plt.xlabel('Time (s)') - plt.ylabel(r'$v_{\parallel}/v$') - # plt.yscale('log') - plt.tight_layout() - plt.savefig(os.path.join(output_dir,f'guiding_center_vpar_SIMSOPT_vs_ESSOS_tolerance{trace_tolerance_SIMSOPT_array[i]}.pdf'), dpi=150) - plt.close() - -# Calculate RMS error for each tolerance -rms_error_array = jnp.array([[jnp.sqrt(jnp.mean(jnp.square(jnp.array(error)))) for error in relative_error] for relative_error in relative_error_array]) - -# Plot RMS error in a bar chart -plt.figure() -bar_width = 0.15 -x = jnp.arange(len(trace_tolerance_SIMSOPT_array)) -for i in range(rms_error_array.shape[1]): - plt.bar(x + i * bar_width, rms_error_array[:, i], bar_width, label=f'Particle {1+i}') -plt.xlabel('Tracing Tolerance of SIMSOPT') -plt.ylabel('RMS Error') -plt.yscale('log') -plt.xticks(x + bar_width * (rms_error_array.shape[1] - 1) / 2, [f'Tol={tol}' for tol in trace_tolerance_SIMSOPT_array], rotation=45) -plt.legend() -plt.tight_layout() -plt.savefig(os.path.join(output_dir, 'rms_error_guiding_center_SIMSOPT_vs_ESSOS.pdf'), dpi=150) -plt.close() - -# Calculate maximum error for each tolerance -max_error_array = jnp.array([[jnp.max(jnp.array(error)) for error in relative_error] for relative_error in relative_error_array]) -# Plot maximum error in a bar chart -plt.figure() -bar_width = 0.15 -x = jnp.arange(len(trace_tolerance_SIMSOPT_array)) -for i in range(max_error_array.shape[1]): - plt.bar(x + i * bar_width, max_error_array[:, i], bar_width, label=f'Particle {1+i}') -plt.xlabel('Tracing Tolerance of SIMSOPT') -plt.ylabel('Maximum Error') -plt.yscale('log') -plt.xticks(x + bar_width * (max_error_array.shape[1] - 1) / 2, [f'Tol={tol}' for tol in trace_tolerance_SIMSOPT_array], rotation=45) -plt.legend() -plt.tight_layout() -plt.savefig(os.path.join(output_dir, 'max_error_guiding_center_SIMSOPT_vs_ESSOS.pdf'), dpi=150) -plt.close() - -# Calculate mean error for each tolerance -mean_error_array = jnp.array([[jnp.mean(jnp.array(error)) for error in relative_error] for relative_error in relative_error_array]) -# Plot mean error in a bar chart -plt.figure() -bar_width = 0.15 -x = jnp.arange(len(trace_tolerance_SIMSOPT_array)) -for i in range(mean_error_array.shape[1]): - plt.bar(x + i * bar_width, mean_error_array[:, i], bar_width, label=f'Particle {1+i}') -plt.xlabel('Tracing Tolerance of SIMSOPT') -plt.ylabel('Mean Error') -plt.yscale('log') -plt.xticks(x + bar_width * (mean_error_array.shape[1] - 1) / 2, [f'Tol={tol}' for tol in trace_tolerance_SIMSOPT_array], rotation=45) -plt.legend() -plt.tight_layout() -plt.savefig(os.path.join(output_dir, 'mean_error_guiding_center_SIMSOPT_vs_ESSOS.pdf'), dpi=150) -plt.close() \ No newline at end of file diff --git a/examples/comparisons_SIMSOPT/surfaces_SIMSOPT_vs_ESSOS.py b/examples/comparisons_SIMSOPT/surfaces_SIMSOPT_vs_ESSOS.py deleted file mode 100644 index 7e1780b3..00000000 --- a/examples/comparisons_SIMSOPT/surfaces_SIMSOPT_vs_ESSOS.py +++ /dev/null @@ -1,72 +0,0 @@ -import os -from time import time -import matplotlib.pyplot as plt -from jax import vmap -import jax.numpy as jnp -from essos.coils import Coils, CreateEquallySpacedCurves -from essos.fields import Vmec, BiotSavart -from essos.surfaces import B_on_surface, BdotN_over_B, SurfaceRZFourier as SurfaceRZFourier_ESSOS -from simsopt.field import BiotSavart as BiotSavart_simsopt -from simsopt.geo import SurfaceRZFourier as SurfaceRZFourier_SIMSOPT -from simsopt.objectives import SquaredFlux - -# Optimization parameters -max_coil_length = 42 -order_Fourier_series_coils = 4 -number_coil_points = 50 -function_evaluations_array = [30]*1 -diff_step_array = [1e-2]*1 -number_coils_per_half_field_period = 3 - -ntheta = 36 -nphi = 32 - -# Initialize VMEC field -vmec_file = os.path.join(os.path.dirname(__file__), '..', 'input_files', - 'wout_LandremanPaul2021_QA_reactorScale_lowres.nc') -vmec = Vmec(vmec_file, ntheta=ntheta, nphi=nphi, close=False) - -# Initialize coils -current_on_each_coil = 1 -number_of_field_periods = vmec.nfp -major_radius_coils = vmec.r_axis -minor_radius_coils = vmec.r_axis/1.5 -curves_essos = CreateEquallySpacedCurves(n_curves=number_coils_per_half_field_period, - order=order_Fourier_series_coils, - R=major_radius_coils, r=minor_radius_coils, - n_segments=number_coil_points, - nfp=number_of_field_periods, stellsym=True) -coils_essos = Coils(curves=curves_essos, currents=[current_on_each_coil]*number_coils_per_half_field_period) -field_essos = BiotSavart(coils_essos) -surface_essos = SurfaceRZFourier_ESSOS(vmec, ntheta=ntheta, nphi=nphi, close=False) -# surface_essos.to_vtk("essos_surface") - -coils_simsopt = coils_essos.to_simsopt() -curves_simsopt = curves_essos.to_simsopt() -field_simsopt = BiotSavart_simsopt(coils_simsopt) -surface_simsopt = SurfaceRZFourier_SIMSOPT.from_wout(vmec_file, range="full torus", nphi=nphi, ntheta=ntheta) -field_simsopt.set_points(surface_simsopt.gamma().reshape((-1, 3))) -# surface_simsopt.to_vtk("simsopt_surface") - -print("Gamma") -print(jnp.sum(jnp.abs(surface_simsopt.gamma()-surface_essos.gamma))) - -print('Gamma dash theta') -print(jnp.sum(jnp.abs(surface_simsopt.gammadash2()-surface_essos.gammadash_theta))) - -print('Gamma dash phi') -print(jnp.sum(jnp.abs(surface_simsopt.gammadash1()-surface_essos.gammadash_phi))) - -print('Normal') -print(jnp.sum(jnp.abs(surface_simsopt.normal()-surface_essos.normal))) - -print('Unit normal') -print(jnp.sum(jnp.abs(surface_simsopt.unitnormal()-surface_essos.unitnormal))) - -BdotN_over_B_SIMSOPT = SquaredFlux(surface_simsopt, field_simsopt, definition="normalized").J() -BdotN_over_B_ESSOS = BdotN_over_B(surface_essos, field_essos) - -B_on_surface_simsopt = field_simsopt.B().reshape(surface_simsopt.normal().shape) -B_on_surface_ESSOS = B_on_surface(surface_essos, field_essos) -# print("ESSOS: ", BdotN_over_B_ESSOS) -# print("SIMSOPT: ", BdotN_over_B_SIMSOPT) diff --git a/examples/comparisons_simsopt/coils.py b/examples/comparisons_simsopt/coils.py new file mode 100644 index 00000000..efb7017f --- /dev/null +++ b/examples/comparisons_simsopt/coils.py @@ -0,0 +1,232 @@ +import os +from time import time +import jax.numpy as jnp +import matplotlib.pyplot as plt +plt.rcParams.update({'font.size': 18}) +from jax import block_until_ready +from essos.fields import BiotSavart as BiotSavart_essos +from essos.coils import Coils, Curves +from simsopt import load +from simsopt.geo import CurveXYZFourier, curves_to_vtk +from simsopt.field import BiotSavart as BiotSavart_simsopt, coils_via_symmetries +from simsopt.configs import get_ncsx_data, get_w7x_data, get_hsx_data, get_giuliani_data + +output_dir = os.path.join(os.path.dirname(__file__), '../output') +if not os.path.exists(output_dir): + os.makedirs(output_dir) + +n_segments = 100 + +LandremanPaulQA_json_file = os.path.join(os.path.dirname(__file__), '../../examples/', 'input_files', 'SIMSOPT_biot_savart_LandremanPaulQA.json') +nfp_array = [3, 2, 5, 4, 2] +curves_array = [get_ncsx_data()[0], LandremanPaulQA_json_file, get_w7x_data()[0], get_hsx_data()[0], get_giuliani_data()[0]] +currents_array = [get_ncsx_data()[1], None, get_w7x_data()[1], get_hsx_data()[1], get_giuliani_data()[1]] +name_array = ["NCSX", "QA(json)", "W7-X", "HSX", "Giuliani"] + +print(f'Output being saved to {output_dir}') +print(f'SIMSOPT LandremanPaulQA json file location: {LandremanPaulQA_json_file}') +for nfp, curves_stel, currents_stel, name in zip(nfp_array, curves_array, currents_array, name_array): + print(f' Running {name} and saving to output directory...') + if currents_stel is None: + json_file_stel = curves_stel + field_simsopt = load(json_file_stel) + coils_simsopt = field_simsopt.coils + curves_simsopt = [coil.curve for coil in coils_simsopt] + currents_simsopt = [coil.current for coil in coils_simsopt] + coils_essos = Coils.from_simsopt(json_file_stel, nfp) + curves_essos = Curves.from_simsopt(json_file_stel, nfp) + else: + coils_simsopt = coils_via_symmetries(curves_stel, currents_stel, nfp, True) + curves_simsopt = [c.curve for c in coils_simsopt] + currents_simsopt = [c.current for c in coils_simsopt] + field_simsopt = BiotSavart_simsopt(coils_simsopt) + + coils_essos = Coils.from_simsopt(coils_simsopt, nfp) + curves_essos = Curves.from_simsopt(curves_simsopt, nfp) + + field_essos = BiotSavart_essos(coils_essos) + + coils_essos_to_simsopt = coils_essos.to_simsopt() + curves_essos_to_simsopt = curves_essos.to_simsopt() + field_essos_to_simsopt = BiotSavart_simsopt(coils_essos_to_simsopt) + + # curves_to_vtk(curves_simsopt, os.path.join(output_dir,f"curves_simsopt_{name}")) + # curves_essos.to_vtk(os.path.join(output_dir,f"curves_essos_{name}")) + # curves_to_vtk(curves_essos_to_simsopt, os.path.join(output_dir,f"curves_essos_to_simsopt_{name}")) + + base_coils_simsopt = coils_simsopt[:int(len(coils_simsopt)/2/nfp)] + R = jnp.mean(jnp.array([jnp.sqrt(coil.curve.x[coil.curve.local_dof_names.index('xc(0)')]**2 + +coil.curve.x[coil.curve.local_dof_names.index('yc(0)')]**2) + for coil in base_coils_simsopt])) + x = jnp.array([R+0.01,R,R]) + y = jnp.array([R,R+0.01,R-0.01]) + z = jnp.array([0.05,0.06,0.07]) + + positions = jnp.array((x,y,z)) + + def update_nsegments_simsopt(curve_simsopt, n_segments): + new_curve = CurveXYZFourier(n_segments, curve_simsopt.order) + new_curve.x = curve_simsopt.x + return new_curve + + coils_essos.n_segments = n_segments + + base_curves_simsopt = [update_nsegments_simsopt(coil_simsopt.curve, n_segments) for coil_simsopt in base_coils_simsopt] + coils_simsopt = coils_via_symmetries(base_curves_simsopt, currents_simsopt[0:len(base_coils_simsopt)], nfp, True) + curves_simsopt = [c.curve for c in coils_simsopt] + + # Running the first time for compilation + [curve.gamma() for curve in curves_simsopt] + [curve.gammadash() for curve in curves_simsopt] + [curve.gammadashdash() for curve in curves_simsopt] + coils_essos.gamma + coils_essos.gamma_dash + coils_essos.gamma_dashdash + coils_essos.curvature + coils_essos.reset_cache() + + # Running the second time for coils characteristics comparison + start_time = time() + gamma_curves_simsopt = block_until_ready(jnp.array([curve.gamma() for curve in curves_simsopt])) + t_gamma_avg_simsopt = time() - start_time + + start_time = time() + gamma_curves_essos = block_until_ready(jnp.array(coils_essos.gamma)) + t_gamma_avg_essos = time() - start_time + + start_time = time() + gammadash_curves_simsopt = block_until_ready(jnp.array([curve.gammadash() for curve in curves_simsopt])) + t_gammadash_avg_simsopt = time() - start_time + + start_time = time() + gammadash_curves_essos = block_until_ready(jnp.array(coils_essos.gamma_dash)) + t_gammadash_avg_essos = time() - start_time + + start_time = time() + gammadashdash_curves_simsopt = block_until_ready(jnp.array([curve.gammadashdash() for curve in curves_simsopt])) + t_gammadashdash_avg_simsopt = time() - start_time + + start_time = time() + gammadashdash_curves_essos = block_until_ready(jnp.array(coils_essos.gamma_dashdash)) + t_gammadashdash_avg_essos = time() - start_time + + start_time = time() + curvature_curves_simsopt = block_until_ready(jnp.array([curve.kappa() for curve in curves_simsopt])) + t_curvature_avg_simsopt = time() - start_time + + start_time = time() + curvature_curves_essos = block_until_ready(jnp.array(coils_essos.curvature)) + t_curvature_avg_essos = time() - start_time + + gamma_error_avg = jnp.linalg.norm(gamma_curves_essos - gamma_curves_simsopt) + gammadash_error_avg = jnp.linalg.norm(gammadash_curves_essos - gammadash_curves_simsopt) + gammadashdash_error_avg = jnp.linalg.norm(gammadashdash_curves_essos - gammadashdash_curves_simsopt) + curvature_error_avg = jnp.linalg.norm(curvature_curves_essos - curvature_curves_simsopt) + + # Magnetic field comparison + + field_essos = BiotSavart_essos(coils_essos) + field_simsopt = BiotSavart_simsopt(coils_simsopt) + + t_B_avg_essos = 0 + t_B_avg_simsopt = 0 + B_error_avg = 0 + t_dB_by_dX_avg_essos = 0 + t_dB_by_dX_avg_simsopt = 0 + dB_by_dX_error_avg = 0 + + for position in positions: + field_essos.B(position) + time1 = time() + result_B_essos = field_essos.B(position) + t_B_avg_essos = t_B_avg_essos + time() - time1 + normB_essos = jnp.linalg.norm(result_B_essos) + + field_simsopt.set_points(jnp.array([position])) + field_simsopt.B() + time3 = time() + field_simsopt.set_points(jnp.array([position])) + result_simsopt = field_simsopt.B() + t_B_avg_simsopt = t_B_avg_simsopt + time() - time3 + normB_simsopt = jnp.linalg.norm(jnp.array(result_simsopt)) + + B_error_avg = B_error_avg + jnp.abs(normB_essos - normB_simsopt) + + field_essos.dB_by_dX(position) + time1 = time() + field_simsopt.set_points(jnp.array([position])) + result_dB_by_dX_essos = field_essos.dB_by_dX(position) + t_dB_by_dX_avg_essos = t_dB_by_dX_avg_essos + time() - time1 + norm_dB_by_dX_essos = jnp.linalg.norm(result_dB_by_dX_essos) + + field_simsopt.dB_by_dX() + time3 = time() + field_simsopt.set_points(jnp.array([position])) + result_dB_by_dX_simsopt = field_simsopt.dB_by_dX() + t_dB_by_dX_avg_simsopt = t_dB_by_dX_avg_simsopt + time() - time3 + norm_dB_by_dX_simsopt = jnp.linalg.norm(jnp.array(result_dB_by_dX_simsopt)) + + dB_by_dX_error_avg = dB_by_dX_error_avg + jnp.abs(norm_dB_by_dX_essos - norm_dB_by_dX_simsopt) + + # Labels and corresponding absolute errors (ESSOS - SIMSOPT) + quantities_errors = [ + (r"$B$", jnp.abs(B_error_avg)), + (r"$B'$", jnp.abs(dB_by_dX_error_avg)), + (r"$\Gamma$", jnp.abs(gamma_error_avg)), + (r"$\Gamma'$", jnp.abs(gammadash_error_avg)), + (r"$\Gamma''$", jnp.abs(gammadashdash_error_avg)), + (r"$\kappa$", jnp.abs(curvature_error_avg)), + ] + + labels = [q[0] for q in quantities_errors] + error_vals = [q[1] for q in quantities_errors] + + X_axis = jnp.arange(len(labels)) + bar_width = 0.6 + + fig, ax = plt.subplots(figsize=(9, 6)) + ax.bar(X_axis, error_vals, bar_width, color="darkorange", edgecolor="black") + + ax.set_xticks(X_axis) + ax.set_xticklabels(labels) + ax.set_ylabel("Absolute error") + ax.set_yscale("log") + ax.set_ylim(1e-17, 1e-12) + ax.grid(axis='y', which='both', linestyle='--', linewidth=0.6) + + plt.tight_layout() + plt.savefig(os.path.join(output_dir, f"comparisons_coils_error_{name}.pdf"), transparent=True) + plt.close() + + + # Labels and corresponding timings + quantities = [ + (r"$B$", t_B_avg_essos, t_B_avg_simsopt), + (r"$B'$", t_dB_by_dX_avg_essos, t_dB_by_dX_avg_simsopt), + (r"$\Gamma$", t_gamma_avg_essos, t_gamma_avg_simsopt), + (r"$\Gamma'$", t_gammadash_avg_essos, t_gammadash_avg_simsopt), + (r"$\Gamma''$", t_gammadashdash_avg_essos, t_gammadashdash_avg_simsopt), + (r"$\kappa$", t_curvature_avg_essos, t_curvature_avg_simsopt), + ] + + labels = [q[0] for q in quantities] + essos_vals = [q[1] for q in quantities] + simsopt_vals = [q[2] for q in quantities] + + X_axis = jnp.arange(len(labels)) + bar_width = 0.35 + + fig, ax = plt.subplots(figsize=(9, 6)) + ax.bar(X_axis - bar_width/2, essos_vals, bar_width, label="ESSOS", color="red", edgecolor="black") + ax.bar(X_axis + bar_width/2, simsopt_vals, bar_width, label="SIMSOPT", color="blue", edgecolor="black") + + ax.set_xticks(X_axis) + ax.set_xticklabels(labels) + ax.set_ylabel("Computation time (s)") + ax.set_yscale("log") + ax.set_ylim(1e-5, 1e-1) + ax.grid(axis='y', which='both', linestyle='--', linewidth=0.6) + ax.legend(fontsize=12) + plt.tight_layout() + plt.savefig(os.path.join(output_dir, f"comparisons_coils_time_{name}.pdf"), transparent=True) + plt.close() diff --git a/examples/comparisons_simsopt/field_lines.py b/examples/comparisons_simsopt/field_lines.py new file mode 100644 index 00000000..33d9ba15 --- /dev/null +++ b/examples/comparisons_simsopt/field_lines.py @@ -0,0 +1,183 @@ +import os +number_of_processors_to_use = 1 # Parallelization, this should divide nparticles +os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' +from time import time +import jax.numpy as jnp +from jax import block_until_ready, random +from simsopt import load +from simsopt.field import (particles_to_vtk, compute_fieldlines, plot_poincare_data) +from essos.coils import Coils +from essos.constants import PROTON_MASS, ONE_EV +from essos.dynamics import Tracing, Particles +from essos.fields import BiotSavart as BiotSavart_essos +import matplotlib.pyplot as plt +from matplotlib.lines import Line2D +plt.rcParams.update({'font.size': 18}) + +tmax_fl = 2000 +nfieldlines = 5 +axis_shift=0.02 +R0 = jnp.linspace(1.2125346+axis_shift, 1.295-axis_shift, nfieldlines) +trace_tolerance_array = [1e-5, 1e-7, 1e-9, 1e-11, 1e-13] +trace_tolerance_ESSOS = 1e-9 +mass=PROTON_MASS +energy=5000*ONE_EV + +output_dir = os.path.join(os.path.dirname(__file__), '../output') +if not os.path.exists(output_dir): + os.makedirs(output_dir) + +nfp=2 +LandremanPaulQA_json_file = os.path.join(os.path.dirname(__file__), '../../examples', 'input_files', 'SIMSOPT_biot_savart_LandremanPaulQA.json') +field_simsopt = load(LandremanPaulQA_json_file) +field_essos = BiotSavart_essos(Coils.from_simsopt(LandremanPaulQA_json_file, nfp)) + +Z0 = jnp.zeros(nfieldlines) +phi0 = jnp.zeros(nfieldlines) +initial_xyz=jnp.array([R0*jnp.cos(phi0), R0*jnp.sin(phi0), Z0]).T +initial_vparallel_over_v = random.uniform(random.PRNGKey(42), (nfieldlines,), minval=-1, maxval=1) + +phis_poincare = [(i/4)*(2*jnp.pi/nfp) for i in range(4)] + +particles = Particles(initial_xyz=initial_xyz, initial_vparallel_over_v=initial_vparallel_over_v, mass=mass, energy=energy) + +# Trace in SIMSOPT +runtime_SIMSOPT_array = [] +trajectories_SIMSOPT_array = [] +avg_steps_SIMSOPT_array = [] + +print(f'Output being saved to {output_dir}') +print(f'SIMSOPT LandremanPaulQA json file location: {LandremanPaulQA_json_file}\n') +for trace_tolerance_SIMSOPT in trace_tolerance_array: + print(f'Tracing SIMSOPT field lines with tolerance={trace_tolerance_SIMSOPT}') + t1 = time() + trajectories_SIMSOPT, trajectories_SIMSOPT_phi_hits = block_until_ready(compute_fieldlines( + field_simsopt, R0, Z0, tmax=tmax_fl, tol=trace_tolerance_SIMSOPT, phis=phis_poincare)) + runtime_SIMSOPT = time() - t1 + runtime_SIMSOPT_array.append(runtime_SIMSOPT) + avg_steps_SIMSOPT = sum([len(l) for l in trajectories_SIMSOPT]) // nfieldlines + avg_steps_SIMSOPT_array.append(avg_steps_SIMSOPT) + # print(trajectories_SIMSOPT_this_tolerance[0].shape) + print(f"Time for SIMSOPT tracing={runtime_SIMSOPT:.3f}s. Avg num steps={avg_steps_SIMSOPT}\n") + trajectories_SIMSOPT_array.append(trajectories_SIMSOPT) + + # particles_to_vtk(trajectories_SIMSOPT_this_tolerance, os.path.join(output_dir,f'guiding_center_SIMSOPT')) + +# Trace in ESSOS +runtime_ESSOS_array = [] +times_essos_array = [] +trajectories_ESSOS_array = [] +relative_energy_error_ESSOS_array = [] + +# Creating a tracing object for compilation +compile_tracing = Tracing(field=field_essos, model='FieldLine', + initial_conditions=jnp.array([R0*jnp.cos(phi0), R0*jnp.sin(phi0), Z0]).T, + maxtime=tmax_fl, timestep=tmax_fl/100, times_to_trace=100, + atol=trace_tolerance_array[0], rtol=trace_tolerance_array[0]) +block_until_ready(compile_tracing.trajectories) + +for index, trace_tolerance_ESSOS in enumerate(trace_tolerance_array): + num_steps_essos = int(avg_steps_SIMSOPT_array[index]) + print(f'Tracing ESSOS field lines with tolerance={trace_tolerance_ESSOS}') + start_time = time() + tracing = Tracing(field=field_essos, model='FieldLine', + initial_conditions=jnp.array([R0*jnp.cos(phi0), R0*jnp.sin(phi0), Z0]).T, + maxtime=tmax_fl, timestep=tmax_fl/num_steps_essos, times_to_trace=num_steps_essos, + atol=trace_tolerance_ESSOS, rtol=trace_tolerance_ESSOS) + block_until_ready(tracing.trajectories) + runtime_ESSOS = time() - start_time + runtime_ESSOS_array.append(runtime_ESSOS) + times_essos_array.append(tracing.times) + trajectories_ESSOS_array.append(tracing.trajectories) + # print(tracing.trajectories.shape) + + trajectories_ESSOS = tracing.trajectories + print(f"Time for ESSOS tracing={runtime_ESSOS:.3f}s. Num steps={len(trajectories_ESSOS[0])}\n") + + # tracing.to_vtk(os.path.join(output_dir,f'guiding_center_ESSOS')) + +print('Plotting the results to output directory...') + +# Plot time comparison in a bar chart +quantities = [(fr"tol=$10^{{{int(jnp.log10(trace_tolerance_array[tolerance_idx])-1e-3)}}}$", runtime_ESSOS_array[tolerance_idx], runtime_SIMSOPT_array[tolerance_idx]) + for tolerance_idx in range(len(trace_tolerance_array))] + +labels = [q[0] for q in quantities] +essos_vals = [q[1] for q in quantities] +simsopt_vals = [q[2] for q in quantities] + +X_axis = jnp.arange(len(labels)) +bar_width = 0.35 + +fig, ax = plt.subplots(figsize=(9, 6)) +ax.bar(X_axis - bar_width/2, essos_vals, bar_width, label="ESSOS", color="red", edgecolor="black") +ax.bar(X_axis + bar_width/2, simsopt_vals, bar_width, label="SIMSOPT", color="blue", edgecolor="black") + +ax.set_xticks(X_axis) +ax.set_xticklabels(labels) +ax.set_ylabel("Computation time (s)") +ax.set_yscale('log') +ax.set_ylim(1e0, 1e2) +ax.grid(axis='y', which='both', linestyle='--', linewidth=0.6) +ax.legend(fontsize=14) +plt.savefig(os.path.join(output_dir, 'comparisons_fl_times.pdf'), dpi=150) + +################################## + +def interpolate_SIMSOPT_to_ESSOS(trajectory_SIMSOPT, time_ESSOS): + time_simsopt = trajectory_SIMSOPT[:, 0] # Time values from SIMSOPT trajectory + + interp_x = jnp.interp(time_ESSOS, time_simsopt, trajectory_SIMSOPT[:, 1]) + interp_y = jnp.interp(time_ESSOS, time_simsopt, trajectory_SIMSOPT[:, 2]) + interp_z = jnp.interp(time_ESSOS, time_simsopt, trajectory_SIMSOPT[:, 3]) + + coords_SIMSOPT_interp = jnp.column_stack([interp_x, interp_y, interp_z]) + + return coords_SIMSOPT_interp + +plt.figure(figsize=(9, 6)) + +avg_relative_xyz_error_array = [] +for tolerance_idx in range(len(trace_tolerance_array)): + this_trajectory_SIMSOPT = jnp.stack([interpolate_SIMSOPT_to_ESSOS( + trajectories_SIMSOPT_array[tolerance_idx][particle_idx], times_essos_array[tolerance_idx] + ) for particle_idx in range(nfieldlines)]) + + this_trajectory_ESSOS = trajectories_ESSOS_array[tolerance_idx] + + relative_xyz_errors = jnp.linalg.norm(this_trajectory_ESSOS[:, :, :3] - this_trajectory_SIMSOPT[:, :, :3], axis=2) / (jnp.linalg.norm(this_trajectory_SIMSOPT[:, :, :3], axis=2) + 1e-12) + + avg_relative_xyz_errors = jnp.mean(relative_xyz_errors, axis=0) + avg_relative_xyz_error_array.append(jnp.mean(avg_relative_xyz_errors)) + + plt.plot(times_essos_array[tolerance_idx], avg_relative_xyz_errors, label=rf'tol=$10^{{{int(jnp.log10(trace_tolerance_array[tolerance_idx])-1e-3)}}}$') + +plt.legend() +plt.xlabel('Time (a.u.)') +plt.yscale('log') + +plt.ylabel(r'Relative $x,y,z$ Error') +plt.savefig(os.path.join(output_dir, f'comparisons_fl_error_xyz.pdf'), dpi=150) + +quantities = [(fr"tol=$10^{{{int(jnp.log10(trace_tolerance_array[tolerance_idx])-1e-3)}}}$", avg_relative_xyz_error_array[tolerance_idx]) + for tolerance_idx in range(len(trace_tolerance_array))] + +labels = [q[0] for q in quantities] +xyz_vals = [q[1] for q in quantities] + +X_axis = jnp.arange(len(labels)) +bar_width = 0.4 + +fig, ax = plt.subplots(figsize=(9, 6)) +ax.bar(X_axis, xyz_vals, bar_width, label=r"x,y,z", color="darkorange", edgecolor="black") + +ax.set_xticks(X_axis) +ax.set_xticklabels(labels) +ax.set_ylabel("Time-averaged relative error") +ax.set_yscale('log') +ax.set_ylim(1e-6, 1e-1) +ax.grid(axis='y', which='both', linestyle='--', linewidth=0.6) +ax.legend(fontsize=14) +plt.savefig(os.path.join(output_dir, 'comparisons_fl_error.pdf'), dpi=150) + +plt.show() \ No newline at end of file diff --git a/examples/comparisons_simsopt/full_orbit.py b/examples/comparisons_simsopt/full_orbit.py new file mode 100644 index 00000000..3ac55057 --- /dev/null +++ b/examples/comparisons_simsopt/full_orbit.py @@ -0,0 +1,264 @@ +import os +number_of_processors_to_use = 1 # Parallelization, this should divide nparticles +os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' +from time import time +import jax.numpy as jnp +from jax import block_until_ready, random +from simsopt import load +from simsopt.field import (particles_to_vtk, trace_particles, plot_poincare_data) +from essos.coils import Coils +from essos.constants import PROTON_MASS, ONE_EV +from essos.dynamics import Tracing, Particles +from essos.fields import BiotSavart as BiotSavart_essos +import matplotlib.pyplot as plt +from matplotlib.lines import Line2D +plt.rcParams.update({'font.size': 18}) + +######################################################################################## +method = 'Boris' # 'Boris' or 'Dopri5' +######################################################################################## + + +tmax = 5e-5 +nparticles = 5 +axis_shft=0.02 +R0 = jnp.linspace(1.2125346+axis_shft, 1.295-axis_shft, nparticles) +trace_tolerance_array = [1e-5, 1e-7, 1e-9, 1e-11, 1e-13] +trace_tolerance_ESSOS = 1e-9 +mass=PROTON_MASS +energy=5000*ONE_EV + +output_dir = os.path.join(os.path.dirname(__file__), '../output') +if not os.path.exists(output_dir): + os.makedirs(output_dir) + +nfp=2 +LandremanPaulQA_json_file = os.path.join(os.path.dirname(__file__), '../../examples', 'input_files', 'SIMSOPT_biot_savart_LandremanPaulQA.json') +field_simsopt = load(LandremanPaulQA_json_file) +field_essos = BiotSavart_essos(Coils.from_simsopt(LandremanPaulQA_json_file, nfp)) + +Z0 = jnp.zeros(nparticles) +phi0 = jnp.zeros(nparticles) +initial_xyz=jnp.array([R0*jnp.cos(phi0), R0*jnp.sin(phi0), Z0]).T +initial_vparallel_over_v = random.uniform(random.PRNGKey(42), (nparticles,), minval=-1, maxval=1) + +particles = Particles(initial_xyz=initial_xyz, initial_vparallel_over_v=initial_vparallel_over_v, mass=mass, energy=energy, field=field_essos) + +# Trace in SIMSOPT +runtime_SIMSOPT_array = [] +trajectories_SIMSOPT_array = [] +avg_steps_SIMSOPT_array = [] +relative_energy_error_SIMSOPT_array = [] +print(f'Output being saved to {output_dir}') +print(f'SIMSOPT LandremanPaulQA json file location: {LandremanPaulQA_json_file}\n') +for trace_tolerance_SIMSOPT in trace_tolerance_array: + print(f'Tracing SIMSOPT full orbit with tolerance={trace_tolerance_SIMSOPT}') + t1 = time() + trajectories_SIMSOPT, trajectories_SIMSOPT_phi_hits = block_until_ready(trace_particles( + field=field_simsopt, xyz_inits=particles.initial_xyz, mass=particles.mass, + parallel_speeds=particles.initial_vparallel, tmax=tmax, mode='full', + charge=particles.charge, Ekin=particles.energy, tol=trace_tolerance_SIMSOPT)) + runtime_SIMSOPT = time() - t1 + runtime_SIMSOPT_array.append(runtime_SIMSOPT) + avg_steps_SIMSOPT = sum([len(l) for l in trajectories_SIMSOPT]) // nparticles + avg_steps_SIMSOPT_array.append(avg_steps_SIMSOPT) + + print(f"Time for SIMSOPT tracing={runtime_SIMSOPT:.3f}s. Avg num steps={avg_steps_SIMSOPT}\n") + trajectories_SIMSOPT_array.append(trajectories_SIMSOPT) + + relative_energy_SIMSOPT = [jnp.abs(0.5 * mass * jnp.sum(jnp.square(trajectory[:, 4:]), axis=1) - particles.energy) / particles.energy + for trajectory in trajectories_SIMSOPT] + + relative_energy_error_SIMSOPT_array.append(relative_energy_SIMSOPT) + + # particles_to_vtk(trajectories_SIMSOPT_this_tolerance, os.path.join(output_dir,f'guiding_center_SIMSOPT')) + +# Trace in ESSOS +runtime_ESSOS_array = [] +times_essos_array = [] +trajectories_ESSOS_array = [] +relative_energy_error_ESSOS_array = [] + +# Creating a tracing object for compilation +if method == 'Dopri5': + compile_tracing = Tracing(field=field_essos, model='FullOrbit', particles=particles, + maxtime=tmax, timestep=tmax/100, times_to_trace=100, + atol=trace_tolerance_array[0], rtol=trace_tolerance_array[0]) +else: + compile_tracing = Tracing(field=field_essos, model='FullOrbit_Boris', particles=particles, + maxtime=tmax, timestep=tmax/100, times_to_trace=100) + +block_until_ready(compile_tracing.trajectories) + +for tolerance_idx, trace_tolerance_ESSOS in enumerate(trace_tolerance_array): + print(f'Tracing ESSOS full orbit with tolerance={trace_tolerance_ESSOS}') + start_time = time() + if method == 'Dopri5': + num_steps_essos = 10000 + tracing = Tracing(field=field_essos, model='FullOrbit', particles=particles, + maxtime=tmax, timestep=tmax/num_steps_essos, times_to_trace=num_steps_essos, + atol=trace_tolerance_ESSOS, rtol=trace_tolerance_ESSOS) + else: + num_steps_essos = int(avg_steps_SIMSOPT_array[tolerance_idx]*3) + tracing = Tracing(field=field_essos, model='FullOrbit_Boris', particles=particles, + maxtime=tmax, timestep=tmax/num_steps_essos, times_to_trace=num_steps_essos) + + block_until_ready(tracing.trajectories) + runtime_ESSOS = time() - start_time + runtime_ESSOS_array.append(runtime_ESSOS) + times_essos_array.append(tracing.times) + trajectories_ESSOS_array.append(tracing.trajectories) + # print(tracing.trajectories.shape) + + trajectories_ESSOS = tracing.trajectories + print(f"Time for ESSOS tracing={runtime_ESSOS:.3f}s. Num steps={len(trajectories_ESSOS[0])}\n") + + relative_energy_error_ESSOS = jnp.abs(tracing.energy()-particles.energy)/particles.energy + relative_energy_error_ESSOS_array.append(relative_energy_error_ESSOS) + # tracing.to_vtk(os.path.join(output_dir,f'guiding_center_ESSOS')) + +print('Plotting the results to output directory...') +plt.figure(figsize=(9, 6)) +colors = ['blue', 'orange', 'green', 'red', 'purple'] + +SIMSOPT_energy_interp = [] + +for tolerance_idx in range(len(trace_tolerance_array)): + interpolation = jnp.stack([ + jnp.interp(times_essos_array[tolerance_idx], trajectories_SIMSOPT_array[tolerance_idx][particle_idx][:, 0], relative_energy_error_SIMSOPT_array[tolerance_idx][particle_idx]) + for particle_idx in range(nparticles) + ]) # This will have shape (nparticles, len(times_essos_array[tolerance_idx])) + SIMSOPT_energy_interp.append(interpolation) + + plt.plot(times_essos_array[tolerance_idx]*1000, jnp.mean(interpolation, axis=0), '--', color=colors[tolerance_idx]) + plt.plot(times_essos_array[tolerance_idx]*1000, jnp.mean(relative_energy_error_ESSOS_array[tolerance_idx], axis=0), '-', color=colors[tolerance_idx]) + +legend_elements = [Line2D([0], [0], color=colors[tolerance_idx], linestyle='-', label=fr"tol=$10^{{{int(jnp.log10(trace_tolerance_array[tolerance_idx])-1e-3)}}}$") + for tolerance_idx in range(len(trace_tolerance_array))] + +plt.legend(handles=legend_elements, loc='lower right', title='ESSOS (─), SIMSOPT (--)', fontsize=14, title_fontsize=14) +plt.yscale('log') +plt.xlabel('Time (ms)') +plt.ylabel('Average relative energy error') +plt.tight_layout() +if method == 'Dopri5': + plt.savefig(os.path.join(output_dir, f'comparisons_fo_error_energy.pdf'), dpi=150) +else: + plt.savefig(os.path.join(output_dir, f'comparisons_fo_boris_error_energy.pdf'), dpi=150) + +# Plot time comparison in a bar chart + +quantities = [(fr"tol=$10^{{{int(jnp.log10(trace_tolerance_array[tolerance_idx])-1e-3)}}}$", runtime_ESSOS_array[tolerance_idx], runtime_SIMSOPT_array[tolerance_idx]) + for tolerance_idx in range(len(trace_tolerance_array))] + +labels = [q[0] for q in quantities] +essos_vals = [q[1] for q in quantities] +simsopt_vals = [q[2] for q in quantities] + +X_axis = jnp.arange(len(labels)) +bar_width = 0.35 + +fig, ax = plt.subplots(figsize=(9, 6)) +ax.bar(X_axis - bar_width/2, essos_vals, bar_width, label="ESSOS", color="red", edgecolor="black") +ax.bar(X_axis + bar_width/2, simsopt_vals, bar_width, label="SIMSOPT", color="blue", edgecolor="black") + +ax.set_xticks(X_axis) +ax.set_xticklabels(labels) +ax.set_ylabel("Computation time (s)") +ax.set_yscale('log') +if method == 'Dopri5': + ax.set_ylim(1e0, 1e3) +else: + ax.set_ylim(1e-1, 1e3) + +ax.grid(axis='y', which='both', linestyle='--', linewidth=0.6) +ax.legend(fontsize=14) +if method == 'Dopri5': + plt.savefig(os.path.join(output_dir, 'comparisons_fo_times.pdf'), dpi=150) +else: + plt.savefig(os.path.join(output_dir, 'comparisons_fo_boris_times.pdf'), dpi=150) + +################################## + +def interpolate_SIMSOPT_to_ESSOS(trajectory_SIMSOPT, time_ESSOS): + time_simsopt = trajectory_SIMSOPT[:, 0] # Time values from SIMSOPT trajectory + + interp_x = jnp.interp(time_ESSOS, time_simsopt, trajectory_SIMSOPT[:, 1]) + interp_y = jnp.interp(time_ESSOS, time_simsopt, trajectory_SIMSOPT[:, 2]) + interp_z = jnp.interp(time_ESSOS, time_simsopt, trajectory_SIMSOPT[:, 3]) + interp_vx = jnp.interp(time_ESSOS, time_simsopt, trajectory_SIMSOPT[:, 4]) + interp_vy = jnp.interp(time_ESSOS, time_simsopt, trajectory_SIMSOPT[:, 5]) + interp_vz = jnp.interp(time_ESSOS, time_simsopt, trajectory_SIMSOPT[:, 6]) + + coords_SIMSOPT_interp = jnp.column_stack([interp_x, interp_y, interp_z, interp_vx, interp_vy, interp_vz]) + + return coords_SIMSOPT_interp + +xyz_error_fig, xyz_error_ax = plt.subplots(figsize=(9, 6)) +v_error_fig, v_error_ax = plt.subplots(figsize=(9, 6)) + +avg_relative_xyz_error_array = [] +avg_relative_v_error_array = [] +for tolerance_idx in range(len(trace_tolerance_array)): + this_trajectory_SIMSOPT = jnp.stack([interpolate_SIMSOPT_to_ESSOS( + trajectories_SIMSOPT_array[tolerance_idx][particle_idx], times_essos_array[tolerance_idx] + ) for particle_idx in range(nparticles)]) + + this_trajectory_ESSOS = trajectories_ESSOS_array[tolerance_idx] + + relative_xyz_errors = jnp.linalg.norm(this_trajectory_ESSOS[:, :, :3] - this_trajectory_SIMSOPT[:, :, :3], axis=2) / (jnp.linalg.norm(this_trajectory_SIMSOPT[:, :, :3], axis=2) + 1e-12) + relative_v_errors = jnp.linalg.norm(this_trajectory_ESSOS[:, :, 3:] - this_trajectory_SIMSOPT[:, :, 3:], axis=2) / (jnp.linalg.norm(this_trajectory_SIMSOPT[:, :, 3:], axis=2) + 1e-12) + + avg_relative_xyz_errors = jnp.mean(relative_xyz_errors, axis=0) + avg_relative_v_errors = jnp.mean(relative_v_errors, axis=0) + avg_relative_xyz_error_array.append(jnp.mean(avg_relative_xyz_errors)) + avg_relative_v_error_array.append(jnp.mean(avg_relative_v_errors)) + + xyz_error_ax.plot(times_essos_array[tolerance_idx]*1000, avg_relative_xyz_errors, label=rf'tol=$10^{{{int(jnp.log10(trace_tolerance_array[tolerance_idx])-1e-3)}}}$') + v_error_ax.plot(times_essos_array[tolerance_idx]*1000, avg_relative_v_errors, label=rf'tol=$10^{{{int(jnp.log10(trace_tolerance_array[tolerance_idx])-1e-3)}}}$') + +for ax, fig in zip([xyz_error_ax, v_error_ax], [xyz_error_fig, v_error_fig]): + ax.legend() + ax.set_xlabel('Time (ms)') + ax.set_yscale('log') + +xyz_error_ax.set_ylabel(r'Relative $x,y,z$ Error') +v_error_ax.set_ylabel(r'Relative $v_x,v_y,v_z$ Error') +if method == 'Dopri5': + xyz_error_fig.savefig(os.path.join(output_dir, f'comparisons_fo_error_xyz.pdf'), dpi=150) + v_error_fig.savefig(os.path.join(output_dir, f'comparisons_fo_error_v.pdf'), dpi=150) +else: + xyz_error_fig.savefig(os.path.join(output_dir, f'comparisons_fo_boris_error_xyz.pdf'), dpi=150) + v_error_fig.savefig(os.path.join(output_dir, f'comparisons_fo_boris_error_v.pdf'), dpi=150) + +quantities = [(fr"tol=$10^{{{int(jnp.log10(trace_tolerance_array[tolerance_idx])-1e-3)}}}$", avg_relative_xyz_error_array[tolerance_idx], avg_relative_v_error_array[tolerance_idx]) + for tolerance_idx in range(len(trace_tolerance_array))] + +labels = [q[0] for q in quantities] +xyz_vals = [q[1] for q in quantities] +v_vals = [q[2] for q in quantities] + +X_axis = jnp.arange(len(labels)) +bar_width = 0.35 + +fig, ax = plt.subplots(figsize=(9, 6)) +ax.bar(X_axis - bar_width/2, xyz_vals, bar_width, label=r"x,y,z", color="red", edgecolor="black") +ax.bar(X_axis + bar_width/2, v_vals, bar_width, label=r"$v_x,v_y,v_z$", color="blue", edgecolor="black") + +ax.set_xticks(X_axis) +ax.set_xticklabels(labels) +ax.set_ylabel("Time-averaged relative error") +ax.set_yscale('log') +ax.set_ylim(1e-6, 1e1) +if method == 'Dopri5': + ax.set_ylim(1e-8, 1e-1) +else: + ax.set_ylim(1e-4, 1e0) +ax.grid(axis='y', which='both', linestyle='--', linewidth=0.6) +ax.legend(fontsize=14) +if method == 'Dopri5': + plt.savefig(os.path.join(output_dir, 'comparisons_fo_errors.pdf'), dpi=150) +else: + plt.savefig(os.path.join(output_dir, 'comparisons_fo_boris_errors.pdf'), dpi=150) + +plt.show() \ No newline at end of file diff --git a/examples/comparisons_simsopt/guiding_center.py b/examples/comparisons_simsopt/guiding_center.py new file mode 100644 index 00000000..0851f13f --- /dev/null +++ b/examples/comparisons_simsopt/guiding_center.py @@ -0,0 +1,233 @@ +import os +number_of_processors_to_use = 1 # Parallelization, this should divide nparticles +os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' +from time import time +import jax.numpy as jnp +from jax import block_until_ready, random +from simsopt import load +from simsopt.field import (particles_to_vtk, trace_particles, plot_poincare_data) +from essos.coils import Coils +from essos.constants import PROTON_MASS, ONE_EV +from essos.dynamics import Tracing, Particles +from essos.fields import BiotSavart as BiotSavart_essos +import matplotlib.pyplot as plt +from matplotlib.lines import Line2D +plt.rcParams.update({'font.size': 18}) + +tmax_gc = 5e-4 +nparticles = 5 +axis_shft=0.02 +R0 = jnp.linspace(1.2125346+axis_shft, 1.295-axis_shft, nparticles) +trace_tolerance_array = [1e-5, 1e-7, 1e-9, 1e-11, 1e-13] +trace_tolerance_ESSOS = 1e-9 +mass=PROTON_MASS +energy=5000*ONE_EV + +output_dir = os.path.join(os.path.dirname(__file__), '../output') +if not os.path.exists(output_dir): + os.makedirs(output_dir) + +nfp=2 +LandremanPaulQA_json_file = os.path.join(os.path.dirname(__file__), '../../examples', 'input_files', 'SIMSOPT_biot_savart_LandremanPaulQA.json') +field_simsopt = load(LandremanPaulQA_json_file) +field_essos = BiotSavart_essos(Coils.from_simsopt(LandremanPaulQA_json_file, nfp)) + +Z0 = jnp.zeros(nparticles) +phi0 = jnp.zeros(nparticles) +initial_xyz=jnp.array([R0*jnp.cos(phi0), R0*jnp.sin(phi0), Z0]).T +initial_vparallel_over_v = random.uniform(random.PRNGKey(42), (nparticles,), minval=-1, maxval=1) + +phis_poincare = [(i/4)*(2*jnp.pi/nfp) for i in range(4)] + +particles = Particles(initial_xyz=initial_xyz, initial_vparallel_over_v=initial_vparallel_over_v, mass=mass, energy=energy) + +# Trace in SIMSOPT +runtime_SIMSOPT_array = [] +trajectories_SIMSOPT_array = [] +avg_steps_SIMSOPT_array = [] +relative_energy_error_SIMSOPT_array = [] +print(f'Output being saved to {output_dir}') +print(f'SIMSOPT LandremanPaulQA json file location: {LandremanPaulQA_json_file}\n') +for trace_tolerance_SIMSOPT in trace_tolerance_array: + print(f'Tracing SIMSOPT guiding center with tolerance={trace_tolerance_SIMSOPT}') + t1 = time() + trajectories_SIMSOPT, trajectories_SIMSOPT_phi_hits = block_until_ready(trace_particles( + field=field_simsopt, xyz_inits=particles.initial_xyz, mass=particles.mass, + parallel_speeds=particles.initial_vparallel, tmax=tmax_gc, mode='gc_vac', + charge=particles.charge, Ekin=particles.energy, tol=trace_tolerance_SIMSOPT)) + runtime_SIMSOPT = time() - t1 + runtime_SIMSOPT_array.append(runtime_SIMSOPT) + avg_steps_SIMSOPT = sum([len(l) for l in trajectories_SIMSOPT]) // nparticles + avg_steps_SIMSOPT_array.append(avg_steps_SIMSOPT) + # print(trajectories_SIMSOPT_this_tolerance[0].shape) + print(f"Time for SIMSOPT tracing={runtime_SIMSOPT:.3f}s. Avg num steps={avg_steps_SIMSOPT}\n") + trajectories_SIMSOPT_array.append(trajectories_SIMSOPT) + + relative_energy_SIMSOPT = [] + for i, trajectory in enumerate(trajectories_SIMSOPT): + xyz = jnp.asarray(trajectory[:, 1:4]) + vpar = trajectory[:, 4] + field_simsopt.set_points(xyz) + AbsB = field_simsopt.AbsB()[:,0] + mu = (particles.energy - particles.mass*vpar[0]**2/2)/AbsB[0] + relative_energy_SIMSOPT.append(jnp.abs(particles.mass*vpar**2/2+mu*AbsB-particles.energy)/particles.energy) + relative_energy_error_SIMSOPT_array.append(relative_energy_SIMSOPT) + + # particles_to_vtk(trajectories_SIMSOPT_this_tolerance, os.path.join(output_dir,f'guiding_center_SIMSOPT')) + +# Trace in ESSOS +runtime_ESSOS_array = [] +times_essos_array = [] +trajectories_ESSOS_array = [] +relative_energy_error_ESSOS_array = [] + +# Creating a tracing object for compilation +compile_tracing = Tracing(field=field_essos, model='GuidingCenter', particles=particles, + maxtime=tmax_gc, timestep=tmax_gc/100, times_to_trace=100, + atol=trace_tolerance_array[0], rtol=trace_tolerance_array[0]) +block_until_ready(compile_tracing.trajectories) + +for index, trace_tolerance_ESSOS in enumerate(trace_tolerance_array): + num_steps_essos = avg_steps_SIMSOPT_array[index] + print(f'Tracing ESSOS guiding center with tolerance={trace_tolerance_ESSOS}') + start_time = time() + num_steps_int = int(num_steps_essos) + tracing = Tracing(field=field_essos, model='GuidingCenter', particles=particles, + maxtime=tmax_gc, timestep=tmax_gc/num_steps_int, times_to_trace=num_steps_int, + atol=trace_tolerance_ESSOS, rtol=trace_tolerance_ESSOS) + block_until_ready(tracing.trajectories) + runtime_ESSOS = time() - start_time + runtime_ESSOS_array.append(runtime_ESSOS) + times_essos_array.append(tracing.times) + trajectories_ESSOS_array.append(tracing.trajectories) + # print(tracing.trajectories.shape) + + trajectories_ESSOS = tracing.trajectories + print(f"Time for ESSOS tracing={runtime_ESSOS:.3f}s. Num steps={len(trajectories_ESSOS[0])}\n") + + relative_energy_error_ESSOS = jnp.abs(tracing.energy()-particles.energy)/particles.energy + relative_energy_error_ESSOS_array.append(relative_energy_error_ESSOS) + # tracing.to_vtk(os.path.join(output_dir,f'guiding_center_ESSOS')) + +print('Plotting the results to output directory...') +plt.figure(figsize=(9, 6)) +colors = ['blue', 'orange', 'green', 'red', 'purple'] + +SIMSOPT_energy_interp = [] + +for tolerance_idx in range(len(trace_tolerance_array)): + interpolation = jnp.stack([ + jnp.interp(times_essos_array[tolerance_idx], trajectories_SIMSOPT_array[tolerance_idx][particle_idx][:, 0], relative_energy_error_SIMSOPT_array[tolerance_idx][particle_idx]) + for particle_idx in range(nparticles) + ]) # This will have shape (nparticles, len(times_essos_array[tolerance_idx])) + SIMSOPT_energy_interp.append(interpolation) + + plt.plot(times_essos_array[tolerance_idx]*1000, jnp.mean(interpolation, axis=0), '--', color=colors[tolerance_idx]) + plt.plot(times_essos_array[tolerance_idx]*1000, jnp.mean(relative_energy_error_ESSOS_array[tolerance_idx], axis=0), '-', color=colors[tolerance_idx]) + +legend_elements = [Line2D([0], [0], color=colors[tolerance_idx], linestyle='-', label=fr"tol=$10^{{{int(jnp.log10(trace_tolerance_array[tolerance_idx])-1e-3)}}}$") + for tolerance_idx in range(len(trace_tolerance_array))] + +plt.legend(handles=legend_elements, loc='lower right', title='ESSOS (─), SIMSOPT (--)', fontsize=14, title_fontsize=14) +plt.yscale('log') +plt.xlabel('Time (ms)') +plt.ylabel('Average relative energy error') +plt.tight_layout() +plt.savefig(os.path.join(output_dir, f'comparisons_gc_error_energy.pdf'), dpi=150) + +# Plot time comparison in a bar chart + +quantities = [(fr"tol=$10^{{{int(jnp.log10(trace_tolerance_array[tolerance_idx])-1e-3)}}}$", runtime_ESSOS_array[tolerance_idx], runtime_SIMSOPT_array[tolerance_idx]) + for tolerance_idx in range(len(trace_tolerance_array))] + +labels = [q[0] for q in quantities] +essos_vals = [q[1] for q in quantities] +simsopt_vals = [q[2] for q in quantities] + +X_axis = jnp.arange(len(labels)) +bar_width = 0.35 + +fig, ax = plt.subplots(figsize=(9, 6)) +ax.bar(X_axis - bar_width/2, essos_vals, bar_width, label="ESSOS", color="red", edgecolor="black") +ax.bar(X_axis + bar_width/2, simsopt_vals, bar_width, label="SIMSOPT", color="blue", edgecolor="black") + +ax.set_xticks(X_axis) +ax.set_xticklabels(labels) +ax.set_ylabel("Computation time (s)") +ax.set_yscale('log') +ax.set_ylim(1e0, 1e2) +ax.grid(axis='y', which='both', linestyle='--', linewidth=0.6) +ax.legend(fontsize=14) +plt.savefig(os.path.join(output_dir, 'comparisons_gc_times.pdf'), dpi=150) + +################################## + +def interpolate_SIMSOPT_to_ESSOS(trajectory_SIMSOPT, time_ESSOS): + time_simsopt = trajectory_SIMSOPT[:, 0] # Time values from SIMSOPT trajectory + + interp_x = jnp.interp(time_ESSOS, time_simsopt, trajectory_SIMSOPT[:, 1]) + interp_y = jnp.interp(time_ESSOS, time_simsopt, trajectory_SIMSOPT[:, 2]) + interp_z = jnp.interp(time_ESSOS, time_simsopt, trajectory_SIMSOPT[:, 3]) + interp_v = jnp.interp(time_ESSOS, time_simsopt, trajectory_SIMSOPT[:, 4]) + + coords_SIMSOPT_interp = jnp.column_stack([interp_x, interp_y, interp_z, interp_v]) + + return coords_SIMSOPT_interp + +xyz_error_fig, xyz_error_ax = plt.subplots(figsize=(9, 6)) +vpar_error_fig, vpar_error_ax = plt.subplots(figsize=(9, 6)) + +avg_relative_xyz_error_array = [] +avg_relative_v_error_array = [] +for tolerance_idx in range(len(trace_tolerance_array)): + this_trajectory_SIMSOPT = jnp.stack([interpolate_SIMSOPT_to_ESSOS( + trajectories_SIMSOPT_array[tolerance_idx][particle_idx], times_essos_array[tolerance_idx] + ) for particle_idx in range(nparticles)]) + + this_trajectory_ESSOS = trajectories_ESSOS_array[tolerance_idx] + + relative_xyz_errors = jnp.linalg.norm(this_trajectory_ESSOS[:, :, :3] - this_trajectory_SIMSOPT[:, :, :3], axis=2) / (jnp.linalg.norm(this_trajectory_SIMSOPT[:, :, :3], axis=2) + 1e-12) + relative_v_errors = jnp.abs(this_trajectory_SIMSOPT[:, :, 3] - this_trajectory_ESSOS[:, :, 3]) / (jnp.abs(this_trajectory_SIMSOPT[:, :, 3]) + 1e-12) + + avg_relative_xyz_errors = jnp.mean(relative_xyz_errors, axis=0) + avg_relative_v_errors = jnp.mean(relative_v_errors, axis=0) + avg_relative_xyz_error_array.append(jnp.mean(avg_relative_xyz_errors)) + avg_relative_v_error_array.append(jnp.mean(avg_relative_v_errors)) + + xyz_error_ax.plot(times_essos_array[tolerance_idx]*1000, avg_relative_xyz_errors, label=rf'tol=$10^{{{int(jnp.log10(trace_tolerance_array[tolerance_idx])-1e-3)}}}$') + vpar_error_ax.plot(times_essos_array[tolerance_idx]*1000, avg_relative_v_errors, label=rf'tol=$10^{{{int(jnp.log10(trace_tolerance_array[tolerance_idx])-1e-3)}}}$') + +for ax, fig in zip([xyz_error_ax, vpar_error_ax], [xyz_error_fig, vpar_error_fig]): + ax.legend() + ax.set_xlabel('Time (ms)') + ax.set_yscale('log') + +xyz_error_ax.set_ylabel(r'Relative $x,y,z$ Error') +vpar_error_ax.set_ylabel(r'Relative $v_\parallel$ Error') +xyz_error_fig.savefig(os.path.join(output_dir, f'comparisons_gc_error_xyz.pdf'), dpi=150) +vpar_error_fig.savefig(os.path.join(output_dir, f'comparisons_gc_error_vpar.pdf'), dpi=150) + +quantities = [(fr"tol=$10^{{{int(jnp.log10(trace_tolerance_array[tolerance_idx])-1e-3)}}}$", avg_relative_xyz_error_array[tolerance_idx], avg_relative_v_error_array[tolerance_idx]) + for tolerance_idx in range(len(trace_tolerance_array))] + +labels = [q[0] for q in quantities] +xyz_vals = [q[1] for q in quantities] +vpar_vals = [q[2] for q in quantities] + +X_axis = jnp.arange(len(labels)) +bar_width = 0.35 + +fig, ax = plt.subplots(figsize=(9, 6)) +ax.bar(X_axis - bar_width/2, xyz_vals, bar_width, label=r"x,y,z", color="red", edgecolor="black") +ax.bar(X_axis + bar_width/2, vpar_vals, bar_width, label=r"$v_\parallel$", color="blue", edgecolor="black") + +ax.set_xticks(X_axis) +ax.set_xticklabels(labels) +ax.set_ylabel("Time-averaged relative error") +ax.set_yscale('log') +ax.set_ylim(1e-6, 1e-1) +ax.grid(axis='y', which='both', linestyle='--', linewidth=0.6) +ax.legend(fontsize=14) +plt.savefig(os.path.join(output_dir, 'comparisons_gc_error.pdf'), dpi=150) + +plt.show() \ No newline at end of file diff --git a/examples/comparisons_simsopt/losses.py b/examples/comparisons_simsopt/losses.py new file mode 100644 index 00000000..837e3ca4 --- /dev/null +++ b/examples/comparisons_simsopt/losses.py @@ -0,0 +1,197 @@ +import os +from time import perf_counter as time +import jax.numpy as jnp +import matplotlib.pyplot as plt +plt.rcParams.update({'font.size': 18}) +from jax import block_until_ready +from essos.fields import BiotSavart as BiotSavart_essos +from essos.coils import Coils, Curves +from essos.objective_functions import loss_coil_curvature, loss_coil_separation, compute_candidates, loss_coil_length +from simsopt import load +from simsopt.geo import CurveXYZFourier, curves_to_vtk, CurveCurveDistance, LpCurveCurvature, CurveLength +from simsopt.field import BiotSavart as BiotSavart_simsopt, coils_via_symmetries +from simsopt.configs import get_ncsx_data, get_w7x_data, get_hsx_data, get_giuliani_data + +output_dir = os.path.join(os.path.dirname(__file__), '../output') +if not os.path.exists(output_dir): + os.makedirs(output_dir) + +n_segments = 100 + +LandremanPaulQA_json_file = os.path.join(os.path.dirname(__file__), '../../examples/', 'input_files', 'SIMSOPT_biot_savart_LandremanPaulQA.json') +nfp_array = [3, 2, 5, 4, 2] +curves_array = [get_ncsx_data()[0], LandremanPaulQA_json_file, get_w7x_data()[0], get_hsx_data()[0], get_giuliani_data()[0]] +currents_array = [get_ncsx_data()[1], None, get_w7x_data()[1], get_hsx_data()[1], get_giuliani_data()[1]] +name_array = ["NCSX", "QA(json)", "W7-X", "HSX", "Giuliani"] + +print(f'Output being saved to {output_dir}') +print(f'SIMSOPT LandremanPaulQA json file location: {LandremanPaulQA_json_file}') +for nfp, curves_stel, currents_stel, name in zip(nfp_array, curves_array, currents_array, name_array): + print(f' Running {name} and saving to output directory...') + if currents_stel is None: + json_file_stel = curves_stel + field_simsopt = load(json_file_stel) + coils_simsopt = field_simsopt.coils + curves_simsopt = [coil.curve for coil in coils_simsopt] + currents_simsopt = [coil.current for coil in coils_simsopt] + coils_essos = Coils.from_simsopt(json_file_stel, nfp) + curves_essos = Curves.from_simsopt(json_file_stel, nfp) + else: + coils_simsopt = coils_via_symmetries(curves_stel, currents_stel, nfp, True) + curves_simsopt = [c.curve for c in coils_simsopt] + currents_simsopt = [c.current for c in coils_simsopt] + field_simsopt = BiotSavart_simsopt(coils_simsopt) + + coils_essos = Coils.from_simsopt(coils_simsopt, nfp) + curves_essos = Curves.from_simsopt(curves_simsopt, nfp) + + field_essos = BiotSavart_essos(coils_essos) + + coils_essos_to_simsopt = coils_essos.to_simsopt() + curves_essos_to_simsopt = curves_essos.to_simsopt() + field_essos_to_simsopt = BiotSavart_simsopt(coils_essos_to_simsopt) + + # curves_to_vtk(curves_simsopt, os.path.join(output_dir,f"curves_simsopt_{name}")) + # curves_essos.to_vtk(os.path.join(output_dir,f"curves_essos_{name}")) + # curves_to_vtk(curves_essos_to_simsopt, os.path.join(output_dir,f"curves_essos_to_simsopt_{name}")) + + base_coils_simsopt = coils_simsopt[:int(len(coils_simsopt)/2/nfp)] + R = jnp.mean(jnp.array([jnp.sqrt(coil.curve.x[coil.curve.local_dof_names.index('xc(0)')]**2 + +coil.curve.x[coil.curve.local_dof_names.index('yc(0)')]**2) + for coil in base_coils_simsopt])) + x = jnp.array([R+0.01,R,R]) + y = jnp.array([R,R+0.01,R-0.01]) + z = jnp.array([0.05,0.06,0.07]) + + positions = jnp.array((x,y,z)) + + def update_nsegments_simsopt(curve_simsopt, n_segments): + new_curve = CurveXYZFourier(n_segments, curve_simsopt.order) + new_curve.x = curve_simsopt.x + return new_curve + + coils_essos.n_segments = n_segments + + base_curves_simsopt = [update_nsegments_simsopt(coil_simsopt.curve, n_segments) for coil_simsopt in base_coils_simsopt] + coils_simsopt = coils_via_symmetries(base_curves_simsopt, currents_simsopt[0:len(base_coils_simsopt)], nfp, True) + curves_simsopt = [c.curve for c in coils_simsopt] + + # Running the first time for compilation + [LpCurveCurvature(curve, p=2, threshold=0).J() for curve in curves_simsopt] + loss_coil_curvature(coils_essos, 0) + [CurveLength(curve).J() for curve in curves_simsopt] + loss_coil_length(coils_essos, 10) + CurveCurveDistance(curves_simsopt, 0.5).J() + loss_coil_separation(coils_essos, 0.5) + + # Running the second time for losses comparison + + start_time = time() + curvature_loss_simsopt = block_until_ready(2*sum([LpCurveCurvature(curve, p=2, threshold=0).J() for curve in curves_simsopt])) + t_curvature_avg_simsopt = time() - start_time + + start_time = time() + curvature_loss_essos = block_until_ready(jnp.sum(loss_coil_curvature(coils_essos, 0))) + t_curvature_avg_essos = time() - start_time + + start_time = time() + length_loss_simsopt = block_until_ready(sum([(CurveLength(curve).J()/10 - 1)**2 for curve in curves_simsopt])) + t_length_avg_simsopt = time() - start_time + print(f"Length loss SIMSOPT: {length_loss_simsopt}") + + start_time = time() + length_loss_essos = block_until_ready(jnp.sum(loss_coil_length(coils_essos, 10))) + t_length_avg_essos = time() - start_time + print(f"Length loss ESSOS: {length_loss_essos}") + + start_time = time() + separation_loss_simsopt = block_until_ready(CurveCurveDistance(curves_simsopt, 0.5).J()) + t_separation_avg_simsopt = time() - start_time + print(f"Separation loss SIMSOPT: {separation_loss_simsopt}") + + start_time = time() + separation_loss_essos = block_until_ready(loss_coil_separation(coils_essos, 0.5)) + t_separation_avg_essos = time() - start_time + print(f"Separation loss ESSOS: {separation_loss_essos}") + + start_time = time() + ind_separation_loss_simsopt = block_until_ready(CurveCurveDistance(curves_simsopt, 0.5).J()) + t_ind_separation_avg_simsopt = time() - start_time + print(f"Independence separation loss SIMSOPT: {ind_separation_loss_simsopt}") + + start_time = time() + ind_separation_loss_essos = block_until_ready(loss_coil_separation(coils_essos, 0.5, candidates=compute_candidates(coils_essos, 0.5))) + t_ind_separation_avg_essos = time() - start_time + print(f"Independence separation loss ESSOS: {ind_separation_loss_essos}") + + length_error_avg = jnp.linalg.norm(length_loss_essos - length_loss_simsopt) / jnp.linalg.norm(length_loss_simsopt) + if length_error_avg == 0: + length_error_avg = jnp.finfo(jnp.float64).eps + curvature_error_avg = jnp.linalg.norm(curvature_loss_essos - curvature_loss_simsopt) / jnp.linalg.norm(curvature_loss_simsopt) + if curvature_error_avg == 0: + curvature_error_avg = jnp.finfo(jnp.float64).eps + separation_error_avg = jnp.linalg.norm(separation_loss_essos - separation_loss_simsopt) / jnp.linalg.norm(separation_loss_simsopt) + ind_separation_error_avg = jnp.linalg.norm(ind_separation_loss_essos - ind_separation_loss_simsopt) / jnp.linalg.norm(ind_separation_loss_simsopt) + print(f"length_error_avg: {length_error_avg:.2e}") + print(f"curvature_error_avg: {curvature_error_avg:.2e}") + print(f"separation_error_avg: {separation_error_avg:.2e}") + # print(f"ind_separation_error_avg: {ind_separation_error_avg:.2e}") + + # Labels and corresponding absolute errors (ESSOS - SIMSOPT) + quantities_errors = [ + (r"$L_\ell$", jnp.abs(length_error_avg)), + (r"$L_\kappa$", jnp.abs(curvature_error_avg)), + (r"$L_\text{sep}$", jnp.abs(separation_error_avg)), + # (r"$L_\text{sep,ind}$", jnp.abs(ind_separation_error_avg)), + ] + + labels = [q[0] for q in quantities_errors] + error_vals = [q[1] for q in quantities_errors] + + X_axis = jnp.arange(len(labels)) + bar_width = 0.6 + + fig, ax = plt.subplots(figsize=(9, 6)) + ax.bar(X_axis, error_vals, bar_width, color="darkorange", edgecolor="black") + + ax.set_xticks(X_axis) + ax.set_xticklabels(labels) + ax.set_ylabel("Relative error") + ax.set_yscale("log") + ax.set_ylim(1e-16, 1e-1) + ax.grid(axis='y', which='both', linestyle='--', linewidth=0.6) + + plt.tight_layout() + plt.savefig(os.path.join(output_dir, f"comparisons_losses_error_{name}.pdf"), transparent=True) + plt.close() + + + # Labels and corresponding timings + quantities = [ + (r"$L_\ell$", t_length_avg_essos, t_length_avg_simsopt), + (r"$L_\kappa$", t_curvature_avg_essos, t_curvature_avg_simsopt), + (r"$L_\text{sep}$", t_separation_avg_essos, t_separation_avg_simsopt), + # (r"$L_\text{sep,ind}$", t_ind_separation_avg_essos, t_ind_separation_avg_simsopt), + ] + + labels = [q[0] for q in quantities] + essos_vals = [q[1] for q in quantities] + simsopt_vals = [q[2] for q in quantities] + + X_axis = jnp.arange(len(labels)) + bar_width = 0.35 + + fig, ax = plt.subplots(figsize=(9, 6)) + ax.bar(X_axis - bar_width/2, essos_vals, bar_width, label="ESSOS", color="red", edgecolor="black") + ax.bar(X_axis + bar_width/2, simsopt_vals, bar_width, label="SIMSOPT", color="blue", edgecolor="black") + + ax.set_xticks(X_axis) + ax.set_xticklabels(labels) + ax.set_ylabel("Computation time (s)") + ax.set_yscale("log") + ax.set_ylim(1e-4, 1e0) + ax.grid(axis='y', which='both', linestyle='--', linewidth=0.6) + ax.legend(fontsize=12) + plt.tight_layout() + plt.savefig(os.path.join(output_dir, f"comparisons_losses_time_{name}.pdf"), transparent=True) + plt.close() diff --git a/examples/comparisons_simsopt/surfaces.py b/examples/comparisons_simsopt/surfaces.py new file mode 100644 index 00000000..f422c5c0 --- /dev/null +++ b/examples/comparisons_simsopt/surfaces.py @@ -0,0 +1,205 @@ +import os +from time import perf_counter as time +import matplotlib.pyplot as plt +plt.rcParams.update({'font.size': 18}) +import jax.numpy as jnp +from jax import block_until_ready +from essos.coils import Coils, CreateEquallySpacedCurves +from essos.fields import Vmec, BiotSavart +from essos.surfaces import B_on_surface, BdotN_over_B, SurfaceRZFourier as SurfaceRZFourier_ESSOS, SquaredFlux as SquaredFlux_ESSOS +from simsopt.field import BiotSavart as BiotSavart_simsopt +from simsopt.geo import SurfaceRZFourier as SurfaceRZFourier_SIMSOPT +from simsopt.objectives import SquaredFlux as SquaredFlux_SIMSOPT + +output_dir = os.path.join(os.path.dirname(__file__), '../output') +if not os.path.exists(output_dir): + os.makedirs(output_dir) + +# Optimization parameters +max_coil_length = 42 +order_Fourier_series_coils = 4 +number_coil_points = 50 +function_evaluations_array = [30]*1 +diff_step_array = [1e-2]*1 +number_coils_per_half_field_period = 3 + +ntheta = 36 +nphi = 32 + +# Initialize VMEC field +vmec_file = os.path.join(os.path.dirname(__file__), '../../examples', 'input_files', + 'wout_LandremanPaul2021_QA_reactorScale_lowres.nc') +vmec = Vmec(vmec_file, ntheta=ntheta, nphi=nphi, close=False) + +# Initialize coils +current_on_each_coil = 1 +number_of_field_periods = vmec.nfp +major_radius_coils = vmec.r_axis +minor_radius_coils = vmec.r_axis/1.5 +curves_essos = CreateEquallySpacedCurves(n_curves=number_coils_per_half_field_period, + order=order_Fourier_series_coils, + R=major_radius_coils, r=minor_radius_coils, + n_segments=number_coil_points, + nfp=number_of_field_periods, stellsym=True) +coils_essos = Coils(curves=curves_essos, currents=[current_on_each_coil]*number_coils_per_half_field_period) +field_essos = BiotSavart(coils_essos) +surface_essos = SurfaceRZFourier_ESSOS.from_vmec(vmec, ntheta=ntheta, nphi=nphi, close=False) +# surface_essos.to_vtk("essos_surface") + +coils_simsopt = coils_essos.to_simsopt() +curves_simsopt = curves_essos.to_simsopt() +field_simsopt = BiotSavart_simsopt(coils_simsopt) +surface_simsopt = SurfaceRZFourier_SIMSOPT.from_wout(vmec_file, range="full torus", nphi=nphi, ntheta=ntheta) +field_simsopt.set_points(surface_simsopt.gamma().reshape((-1, 3))) +# surface_simsopt.to_vtk("simsopt_surface") + +# Running the first time for compilation +surface_simsopt.gamma() +surface_simsopt.gammadash1() +surface_simsopt.gammadash2() +surface_simsopt.unitnormal() +field_simsopt.B() +SquaredFlux_SIMSOPT(surface_simsopt, field_simsopt).J() +block_until_ready(surface_essos.gamma) + +# Running the second time for surface characteristics comparison + +print("Gamma") +start_time = time() +gamma_essos = block_until_ready(surface_essos.gamma) +t_gamma_essos = time() - start_time + +gamma_simsopt = block_until_ready(surface_simsopt.gamma()) +start_time = time() +t_gamma_simsopt = time() - start_time + +gamma_error = jnp.sum(jnp.abs(gamma_simsopt - gamma_essos)) +print(gamma_error) + + +print('Gamma dash theta') +start_time = time() +gamma_dash_theta_essos = block_until_ready(surface_essos.gammadash_theta) +t_gamma_dash_theta_essos = time() - start_time + +start_time = time() +gamma_dash_theta_simsopt = block_until_ready(surface_simsopt.gammadash2()) +t_gamma_dash_theta_simsopt = time() - start_time + +gamma_dash_theta_error = jnp.sum(jnp.abs(gamma_dash_theta_simsopt - gamma_dash_theta_essos)) +print(gamma_dash_theta_error) + + +print('Gamma dash phi') +start_time = time() +gamma_dash_phi_essos = block_until_ready(surface_essos.gammadash_phi) +t_gamma_dash_phi_essos = time() - start_time + +start_time = time() +gamma_dash_phi_simsopt = block_until_ready(surface_simsopt.gammadash1()) +t_gamma_dash_phi_simsopt = time() - start_time + +gamma_dash_phi_error = jnp.sum(jnp.abs(gamma_dash_phi_simsopt - gamma_dash_phi_essos)) +print(gamma_dash_phi_error) + + +print('Unit normal') +start_time = time() +unit_normal_essos = block_until_ready(surface_essos.unitnormal) +t_unit_normal_essos = time() - start_time + +start_time = time() +unit_normal_simsopt = block_until_ready(surface_simsopt.unitnormal()) +t_unit_normal_simsopt = time() - start_time + +unit_normal_error = jnp.sum(jnp.abs(unit_normal_simsopt - unit_normal_essos)) +print(unit_normal_error) + + +print('B on surface') +start_time = time() +B_on_surface_essos = block_until_ready(B_on_surface(surface_essos, field_essos)) +t_B_on_surface_essos = time() - start_time + +start_time = time() +B_on_surface_simsopt = block_until_ready(field_simsopt.B()) +t_B_on_surface_simsopt = time() - start_time + +B_on_surface_error = jnp.sum(jnp.abs(B_on_surface_simsopt.reshape((nphi, ntheta, 3)) - B_on_surface_essos)) +print(B_on_surface_error) + + +definition = "local" +print("Squared flux", definition) +start_time = time() +sf_essos = block_until_ready(SquaredFlux_ESSOS(surface_essos, field_essos, definition=definition)) +t_squared_flux_essos = time() - start_time + +start_time = time() +sf_simsopt = block_until_ready(SquaredFlux_SIMSOPT(surface_simsopt, field_simsopt, definition=definition).J()) +t_squared_flux_simsopt = time() - start_time + +squared_flux_error = jnp.abs(sf_simsopt - sf_essos) +print(squared_flux_error) + +# Labels and corresponding absolute errors (ESSOS - SIMSOPT) +quantities_errors = [ + (r"$\Gamma$", gamma_error), + (r"$\Gamma'_\theta$", gamma_dash_theta_error), + (r"$\Gamma'_\phi$", gamma_dash_phi_error), + (r"$\mathbf{n}$", unit_normal_error), + # (r"$\mathbf{B}$", B_on_surface_error), + (r"$L_\text{flux}$", squared_flux_error), +] + +labels = [q[0] for q in quantities_errors] +error_vals = [q[1] for q in quantities_errors] + +X_axis = jnp.arange(len(labels)) +bar_width = 0.6 + +fig, ax = plt.subplots(figsize=(9, 6)) +ax.bar(X_axis, error_vals, bar_width, color="darkorange", edgecolor="black") + +ax.set_xticks(X_axis) +ax.set_xticklabels(labels) +ax.set_ylabel("Absolute error") +ax.set_yscale("log") +ax.set_ylim(1e-14, 1e-10) +ax.grid(axis='y', which='both', linestyle='--', linewidth=0.6) + +plt.tight_layout() +plt.savefig(os.path.join(output_dir, f"comparisons_surfaces_error.pdf"), transparent=True) + +# Labels and corresponding timings +quantities = [ + (r"$\Gamma$", t_gamma_essos, t_gamma_simsopt), + (r"$\Gamma'_\theta$", t_gamma_dash_theta_essos, t_gamma_dash_theta_simsopt), + (r"$\Gamma'_\phi$", t_gamma_dash_phi_essos, t_gamma_dash_phi_simsopt), + (r"$\mathbf{n}$", t_unit_normal_essos, t_unit_normal_simsopt), + # (r"$\mathbf{B}$", t_B_on_surface_essos, t_B_on_surface_simsopt), + (r"$L_\text{flux}$", t_squared_flux_essos, t_squared_flux_simsopt), +] + +labels = [q[0] for q in quantities] +essos_vals = [q[1] for q in quantities] +simsopt_vals = [q[2] for q in quantities] + +X_axis = jnp.arange(len(labels)) +bar_width = 0.35 + +fig, ax = plt.subplots(figsize=(9, 6)) +ax.bar(X_axis - bar_width/2, essos_vals, bar_width, label="ESSOS", color="red", edgecolor="black") +ax.bar(X_axis + bar_width/2, simsopt_vals, bar_width, label="SIMSOPT", color="blue", edgecolor="black") + +ax.set_xticks(X_axis) +ax.set_xticklabels(labels) +ax.set_ylabel("Computation time (s)") +ax.set_yscale("log") +ax.set_ylim(1e-7, 1e-1) +ax.grid(axis='y', which='both', linestyle='--', linewidth=0.6) +ax.legend(fontsize=12) +plt.tight_layout() +plt.savefig(os.path.join(output_dir, f"comparisons_surfaces_time.pdf"), transparent=True) + +plt.show() diff --git a/examples/comparisons_SIMSOPT/vmec_SIMSOPT_vs_ESSOS.py b/examples/comparisons_simsopt/vmec_import.py similarity index 50% rename from examples/comparisons_SIMSOPT/vmec_SIMSOPT_vs_ESSOS.py rename to examples/comparisons_simsopt/vmec_import.py index 8c0c1d27..adffffbc 100644 --- a/examples/comparisons_SIMSOPT/vmec_SIMSOPT_vs_ESSOS.py +++ b/examples/comparisons_simsopt/vmec_import.py @@ -2,18 +2,21 @@ from time import time import jax.numpy as jnp import matplotlib.pyplot as plt +plt.rcParams.update({'font.size': 18}) from jax import block_until_ready, random from essos.fields import Vmec as Vmec_essos from simsopt.mhd import Vmec as Vmec_simsopt, vmec_compute_geometry -output_dir = os.path.join(os.path.dirname(__file__), 'output') + +output_dir = os.path.join(os.path.dirname(__file__), '../output') if not os.path.exists(output_dir): os.makedirs(output_dir) -wout_array = [os.path.join(os.path.dirname(__file__), '..', 'input_files', "wout_LandremanPaul2021_QA_reactorScale_lowres.nc"), - os.path.join(os.path.dirname(__file__), '..', 'input_files', "wout_n3are_R7.75B5.7.nc")] +wout_array = [os.path.join(os.path.dirname(__file__), '../../examples/', 'input_files', "wout_LandremanPaul2021_QA_reactorScale_lowres.nc"), + os.path.join(os.path.dirname(__file__), '../../examples/', 'input_files', "wout_n3are_R7.75B5.7.nc")] name_array = ["LandremanPaulQA", 'NCSX'] + print(f'Output being saved to {output_dir}') for name, wout in zip(name_array, wout_array): print(f' Running comparison with VMEC file located at: {wout}') @@ -71,41 +74,66 @@ def timed_B(s, function): average_time_modB_essos /= len(s_array) average_time_B_essos /= len(s_array) average_time_B_simsopt /= len(s_array) - - fig = plt.figure(figsize = (8, 6)) - X_axis = jnp.arange(4) - Y_axis = [average_time_modB_simsopt, average_time_B_simsopt, average_time_modB_essos, average_time_B_essos] - colors = ['blue', 'blue', 'red', 'red'] - hatches = ['/', '\\', '/', '\\'] - bars = plt.bar(X_axis, Y_axis, width=0.4, color=colors) - for bar, hatch in zip(bars, hatches): bar.set_hatch(hatch) - plt.xticks(X_axis, [r"$|\boldsymbol{B}|$ SIMSOPT", r"$\boldsymbol{B}$ SIMSOPT", r"$|\boldsymbol{B}|$ ESSOS", r"$\boldsymbol{B}$ ESSOS"], fontsize=16) - plt.tick_params(axis='both', which='major', labelsize=14) - plt.tick_params(axis='both', which='minor', labelsize=14) - plt.ylabel("Time to evaluate VMEC field (s)", fontsize=14) - plt.grid(axis='y') - plt.yscale("log") - plt.ylim(1e-6, 1) - plt.title(name, fontsize=14) + error_modB /= len(s_array) + error_B /= len(s_array) + + # Labels and corresponding absolute errors (ESSOS - SIMSOPT) + quantities_errors = [ + (r"$B$", jnp.mean(error_modB)), + (r"$\mathbf{B}$", jnp.mean(error_B)), + ] + + labels = [q[0] for q in quantities_errors] + error_vals = [q[1] for q in quantities_errors] + + X_axis = jnp.arange(len(labels)) + bar_width = 0.4 + + fig, ax = plt.subplots(figsize=(9, 6)) + ax.bar(X_axis, error_vals, bar_width, color="darkorange", edgecolor="black") + + ax.set_xticks(X_axis) + ax.set_xticklabels(labels) + ax.set_ylabel("Relative error") + ax.set_yscale("log") + ax.set_ylim(1e-6, 1e-2) + ax.grid(axis='y', which='both', linestyle='--', linewidth=0.6) + plt.tight_layout() - plt.savefig(os.path.join(output_dir,f"time_VMEC_SIMSOPT_vs_ESSOS_{name}.pdf"), transparent=True) - plt.close() - - fig = plt.figure(figsize = (8, 6)) - X_axis = jnp.arange(2) - Y_axis = [jnp.mean(error_modB), jnp.mean(error_B)] - colors = ['purple', 'orange'] - hatches = ['/', '//'] - bars = plt.bar(X_axis, Y_axis, width=0.4, color=colors) - for bar, hatch in zip(bars, hatches): bar.set_hatch(hatch) - plt.xticks(X_axis, [r"$|\boldsymbol{B}|$", r"$\boldsymbol{B}$"], fontsize=16) - plt.tick_params(axis='both', which='major', labelsize=14) - plt.tick_params(axis='both', which='minor', labelsize=14) - plt.ylabel("Relative error SIMSOPT vs ESSOS (%)", fontsize=14) - plt.grid(axis='y') - plt.yscale("log") - plt.ylim(1e-6, 1e-1) - plt.title(name, fontsize=14) + plt.savefig(os.path.join(output_dir, f"comparisons_VMEC_error_{name}.pdf"), transparent=True) + + # Labels and corresponding timings + print(f"Average time to compute |B| in SIMSOPT: {average_time_modB_simsopt:.6f} s") + print(f"Average time to compute B in SIMSOPT: {average_time_B_simsopt:.6f} s") + print(f"Average time to compute |B| in ESSOS: {average_time_modB_essos:.6f} s") + print(f"Average time to compute B in ESSOS: {average_time_B_essos:.6f} s") + print(f"Relative error in |B|: {jnp.mean(error_modB):.6f}") + print(f"Relative error in B: {jnp.mean(error_B):.6f}") + + quantities = [ + (r"$B$", average_time_modB_essos, average_time_modB_simsopt), + (r"$\mathbf{B}$", average_time_B_essos, average_time_B_simsopt), + ] + + labels = [q[0] for q in quantities] + essos_vals = [q[1] for q in quantities] + simsopt_vals = [q[2] for q in quantities] + + X_axis = jnp.arange(len(labels)) + bar_width = 0.4 + + fig, ax = plt.subplots(figsize=(9, 6)) + ax.bar(X_axis - bar_width/2, essos_vals, bar_width, label="ESSOS", color="red", edgecolor="black") + ax.bar(X_axis + bar_width/2, simsopt_vals, bar_width, label="SIMSOPT", color="blue", edgecolor="black") + + ax.set_xticks(X_axis) + ax.set_xticklabels(labels) + ax.set_ylabel("Computation time (s)") + ax.set_yscale("log") + ax.set_ylim(1e-5, 1e-1) + ax.grid(axis='y', which='both', linestyle='--', linewidth=0.6) + ax.legend(fontsize=12) plt.tight_layout() - plt.savefig(os.path.join(output_dir,f"error_VMEC_SIMSOPT_vs_ESSOS_{name}.pdf"), transparent=True) - plt.close() \ No newline at end of file + plt.savefig(os.path.join(output_dir, f"comparisons_VMEC_time_{name}.pdf"), transparent=True) + + plt.show() \ No newline at end of file diff --git a/examples/poincare_guiding_center_coils.py b/examples/fieldline_tracing/poincare_guiding_center_coils.py similarity index 90% rename from examples/poincare_guiding_center_coils.py rename to examples/fieldline_tracing/poincare_guiding_center_coils.py index e0449fdc..e644fca8 100644 --- a/examples/poincare_guiding_center_coils.py +++ b/examples/fieldline_tracing/poincare_guiding_center_coils.py @@ -5,7 +5,7 @@ import jax.numpy as jnp import matplotlib.pyplot as plt from essos.fields import BiotSavart -from essos.coils import Coils_from_json +from essos.coils import Coils from essos.constants import PROTON_MASS, ONE_EV from essos.dynamics import Tracing, Particles @@ -22,8 +22,8 @@ angle = 45 # Load coils and field -json_file = os.path.join(os.path.dirname(__file__), 'input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') -coils = Coils_from_json(json_file) +json_file = os.path.join(os.path.dirname(__file__), '..', 'input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') +coils = Coils.from_json(json_file) field = BiotSavart(coils) # Initialize particles diff --git a/examples/trace_fieldlines_coils.py b/examples/fieldline_tracing/trace_fieldlines_coils.py similarity index 89% rename from examples/trace_fieldlines_coils.py rename to examples/fieldline_tracing/trace_fieldlines_coils.py index 2ea23059..1d66721e 100644 --- a/examples/trace_fieldlines_coils.py +++ b/examples/fieldline_tracing/trace_fieldlines_coils.py @@ -6,7 +6,7 @@ from jax import block_until_ready import matplotlib.pyplot as plt from essos.fields import BiotSavart -from essos.coils import Coils_from_json +from essos.coils import Coils from essos.dynamics import Tracing # Input parameters @@ -18,8 +18,8 @@ num_steps = 6000 # Load coils and field -json_file = os.path.join(os.path.dirname(__file__), 'input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') -coils = Coils_from_json(json_file) +json_file = os.path.join(os.path.dirname(__file__), '..', 'input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') +coils = Coils.from_json(json_file) field = BiotSavart(coils) # Initialize particles diff --git a/examples/trace_fieldlines_vmec.py b/examples/fieldline_tracing/trace_fieldlines_vmec.py similarity index 94% rename from examples/trace_fieldlines_vmec.py rename to examples/fieldline_tracing/trace_fieldlines_vmec.py index fad2ed0e..65eb9962 100644 --- a/examples/trace_fieldlines_vmec.py +++ b/examples/fieldline_tracing/trace_fieldlines_vmec.py @@ -17,7 +17,7 @@ num_steps = 10000 # Load coils and field -wout_file = os.path.join(os.path.dirname(__file__), 'input_files',"wout_QH_simple_scaled.nc") +wout_file = os.path.join(os.path.dirname(__file__), '..', 'input_files',"wout_QH_simple_scaled.nc") vmec = Vmec(wout_file) # Initialize particles diff --git a/examples/get_derivatives_coils_particle_confinement_guidingcenter.py b/examples/get_derivatives_coils_particle_confinement_guidingcenter.py deleted file mode 100644 index 35646d71..00000000 --- a/examples/get_derivatives_coils_particle_confinement_guidingcenter.py +++ /dev/null @@ -1,68 +0,0 @@ - -import os -number_of_processors_to_use = 1 # Parallelization, this should divide nparticles -os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' -from jax import grad -import jax.numpy as jnp -from essos.dynamics import Particles -from essos.coils import Coils, CreateEquallySpacedCurves -from essos.objective_functions import loss_coil_curvature,loss_coil_length,loss_normB_axis_average -from functools import partial - -# Optimization parameters -target_B_on_axis = 5.7 -max_coil_length = 31 -max_coil_curvature = 0.4 -n_particles_per_core=1 -nparticles = number_of_processors_to_use*n_particles_per_core -order_Fourier_series_coils = 2 -number_coil_points = 80 -maximum_function_evaluations = 301 -maxtimes = [1.e-6] -t=maxtimes[0] -num_steps=100 -number_coils_per_half_field_period = 3 -number_of_field_periods = 2 -model = 'GuidingCenterAdaptative' - -# Initialize coils -current_on_each_coil = 1.84e7 -major_radius_coils = 7.75 -minor_radius_coils = 4.45 -curves = CreateEquallySpacedCurves(n_curves=number_coils_per_half_field_period, - order=order_Fourier_series_coils, - R=major_radius_coils, r=minor_radius_coils, - n_segments=number_coil_points, - nfp=number_of_field_periods, stellsym=True) -coils_initial = Coils(curves=curves, currents=[current_on_each_coil]*number_coils_per_half_field_period) - -len_dofs_curves = len(jnp.ravel(coils_initial.dofs_curves)) -nfp = coils_initial.nfp -stellsym = coils_initial.stellsym -n_segments = coils_initial.n_segments -dofs_curves_shape = coils_initial.dofs_curves.shape -currents_scale = coils_initial.currents_scale - -# Initialize particles -phi_array = jnp.linspace(0, 2*jnp.pi, nparticles) -initial_xyz=jnp.array([major_radius_coils*jnp.cos(phi_array), major_radius_coils*jnp.sin(phi_array), 0*phi_array]).T -particles = Particles(initial_xyz=initial_xyz) - -# Objective functions -## Curvature -curvature_partial=partial(loss_coil_curvature, dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,max_coil_curvature=max_coil_curvature) -## Length -length_partial=partial(loss_coil_length, dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,max_coil_length=max_coil_length) -## B on axis -Baxis_average_partial=partial(loss_normB_axis_average,dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,npoints=15,target_B_on_axis=target_B_on_axis) -## All terms put together -def total_loss(params): - return jnp.sum(jnp.square(curvature_partial(params)+length_partial(params)+Baxis_average_partial(params))) - -## Take the gradients -params=coils_initial.x -loss=total_loss(params) -gradients=grad(total_loss)(params) - -print('Objective function: {:.2E}'.format(loss)) -print('Gradients (derivative of objective function with respect to coils): ',gradients) diff --git a/examples/input_files/input.rotating_ellipse b/examples/input_files/input.rotating_ellipse index a35f3af0..bce19ba5 100644 --- a/examples/input_files/input.rotating_ellipse +++ b/examples/input_files/input.rotating_ellipse @@ -5,10 +5,17 @@ MPOL = 002 NTOR = 002 !----- Boundary Parameters (n,m) ----- - RBC( 000,000) = 10 ZBS( 000,000) = 0 - RBC( 001,000) = 1 ZBS( 001,000) = -1 + RBC( 000,000) = 10. ZBS( 000,000) = 0. + RBC( 001,000) = 1. ZBS( 001,000) = -1. + RBC( 002,000) = 0. ZBS( 002,000) = 0. + RBC( -002,001) = 0. ZBS( -002,001) = 0. RBC(-001,001) = 0.1 ZBS(-001,001) = 0.1 RBC( 000,001) = 2.5 ZBS( 000,001) = 2.5 RBC( 001,001) = -1 ZBS( 001,001) = 1 + RBC( 002,001) = 0 ZBS( 002,001) = 0 RBC(-002,002) = 1E-4 ZBS(-002,002) = 1E-4 + RBC(-001,002) = 0. ZBS(-001,002) = 0. + RBC( 000,002) = 0. ZBS( 000,002) = 0. + RBC( 001,002) = 0. ZBS( 001,002) = 0. + RBC( 002,002) = 0. ZBS( 002,002) = 0. / diff --git a/examples/input_files/input.toroidal_surface b/examples/input_files/input.toroidal_surface index 3a133b24..533ae617 100644 --- a/examples/input_files/input.toroidal_surface +++ b/examples/input_files/input.toroidal_surface @@ -1,14 +1,21 @@ !----- Runtime Parameters ----- &INDATA LASYM = F - NFP = 0001 + NFP = 0002 MPOL = 002 NTOR = 002 !----- Boundary Parameters (n,m) ----- - RBC( 000,000) = 7.75 ZBS( 000,000) = 0 + RBC( 000,000) = 10.0 ZBS( 000,000) = 0 RBC( 001,000) = 0.000001 ZBS( 001,000) = -0.000001 + RBC( 002,000) = 0. ZBS( 002,000) = 0. + RBC( -002,001) = 0. ZBS( -002,001) = 0. RBC(-001,001) = 0.000001 ZBS(-001,001) = 0.000001 - RBC( 000,001) = 2.5 ZBS( 000,001) = 2.5 + RBC( 000,001) = 0.5 ZBS( 000,001) = 0.5 RBC( 001,001) = 0.000001 ZBS( 001,001) = 0.000001 + RBC( 002,001) = 0 ZBS( 002,001) = 0 RBC(-002,002) = 1E-7 ZBS(-002,002) = 1E-7 + RBC(-001,002) = 0. ZBS(-001,002) = 0. + RBC( 000,002) = 0. ZBS( 000,002) = 0. + RBC( 001,002) = 0. ZBS( 001,002) = 0. + RBC( 002,002) = 0. ZBS( 002,002) = 0. / diff --git a/examples/optimize_coils_and_nearaxis.py b/examples/optimize_coils_and_nearaxis.py deleted file mode 100644 index 00fa24e4..00000000 --- a/examples/optimize_coils_and_nearaxis.py +++ /dev/null @@ -1,120 +0,0 @@ -import os -number_of_processors_to_use = 4 # Parallelization, this should divide nfieldlines -os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' -from time import time -import jax.numpy as jnp -import matplotlib.pyplot as plt -from essos.coils import Coils, CreateEquallySpacedCurves -from essos.fields import near_axis, BiotSavart -from essos.dynamics import Tracing -from essos.optimization import optimize_loss_function -from essos.objective_functions import (difference_B_gradB_onaxis, - loss_coils_and_nearaxis, loss_coils_for_nearaxis) - -# Optimization parameters -max_coil_length = 4. -max_coil_curvature = 6. -order_Fourier_series_coils = 5 -number_coil_points = order_Fourier_series_coils*10 -maximum_function_evaluations = 200 -number_coils_per_half_field_period = 3 -tolerance_optimization = 1e-8 - -# Initialize Near-Axis field -rc=jnp.array([1, 0.045]) -zs=jnp.array([0,-0.045]) -etabar=-0.9 -nfp=3 -field_nearaxis_initial = near_axis(rc=rc, zs=zs, etabar=etabar, nfp=nfp) - -# Initialize coils -current_on_each_coil = 17e5*field_nearaxis_initial.B0/nfp/2 -number_of_field_periods = nfp -major_radius_coils = field_nearaxis_initial.R0[0] -minor_radius_coils = major_radius_coils/2.0 -curves = CreateEquallySpacedCurves(n_curves=number_coils_per_half_field_period, - order=order_Fourier_series_coils, - R=major_radius_coils, r=minor_radius_coils, - n_segments=number_coil_points, - nfp=number_of_field_periods, stellsym=True) -coils_initial = Coils(curves=curves, currents=[current_on_each_coil]*number_coils_per_half_field_period) - -# Optimize coils -print(f'Optimizing coils for initial near=axis with {maximum_function_evaluations} function evaluations.') -time0 = time() -initial_dofs = coils_initial.x -coils_optimized_initial_nearaxis = optimize_loss_function(loss_coils_for_nearaxis, initial_dofs=coils_initial.x, coils=coils_initial, tolerance_optimization=tolerance_optimization, - maximum_function_evaluations=maximum_function_evaluations, field_nearaxis=field_nearaxis_initial, - max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature,) -print(f"Optimization took {time()-time0:.2f} seconds") - -# Optimize coils -print(f'Optimizing coils and near-axis with {maximum_function_evaluations} function evaluations.') -time0 = time() -initial_dofs = jnp.concatenate((coils_optimized_initial_nearaxis.x, field_nearaxis_initial.x)) -coils_optimized, field_nearaxis_optimized = optimize_loss_function(loss_coils_and_nearaxis, initial_dofs=initial_dofs, coils=coils_initial, tolerance_optimization=tolerance_optimization, - maximum_function_evaluations=maximum_function_evaluations, field_nearaxis=field_nearaxis_initial, - max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature,) -print(f"Optimization took {time()-time0:.2f} seconds") - -B_difference_initial, gradB_difference_initial = difference_B_gradB_onaxis(field_nearaxis_initial, BiotSavart(coils_optimized_initial_nearaxis)) -B_difference_loss_initial = jnp.sum(jnp.abs(B_difference_initial)) -gradB_difference_loss_initial = jnp.sum(jnp.abs(gradB_difference_initial)) - -B_difference_optimized, gradB_difference_optimized = difference_B_gradB_onaxis(field_nearaxis_optimized, BiotSavart(coils_optimized)) -B_difference_loss_optimized = jnp.sum(jnp.abs(B_difference_optimized)) -gradB_difference_loss_optimized = jnp.sum(gradB_difference_optimized) - -print(f'############################################') -print(f'Iota for initial near-axis: {field_nearaxis_initial.iota}') -print(f'Iota for optimized near-axis: {field_nearaxis_optimized.iota}') -print(f'Maximum elongation for initial near-axis: {max(field_nearaxis_initial.elongation)}') -print(f'Maximum elongation for optimized near-axis: {max(field_nearaxis_optimized.elongation)}') -print(f'Loss of B difference for initial near-axis: {B_difference_loss_initial}') -print(f'Loss of B difference for optimized near-axis: {B_difference_loss_optimized}') -print(f'Loss of gradB difference for initial near-axis: {gradB_difference_loss_initial}') -print(f'Loss of gradB difference for optimized near-axis: {gradB_difference_loss_optimized}') - -# Trace fieldlines -nfieldlines = 6 -num_steps = 1000 -tmax = 1.e-6 -trace_tolerance = 1e-7 - -R0_initial = jnp.linspace(field_nearaxis_initial.R0[0], 1.05*field_nearaxis_initial.R0[0], nfieldlines) -R0_optimized = jnp.linspace(field_nearaxis_optimized.R0[0], 1.05*field_nearaxis_optimized.R0[0], nfieldlines) -Z0 = jnp.zeros(nfieldlines) -phi0 = jnp.zeros(nfieldlines) -initial_xyz_initial = jnp.array([R0_initial*jnp.cos(phi0), R0_initial*jnp.sin(phi0), Z0]).T -initial_xyz_optimized = jnp.array([R0_optimized*jnp.cos(phi0), R0_optimized*jnp.sin(phi0), Z0]).T - -time0 = time() -tracing_initial = Tracing(field=BiotSavart(coils_optimized_initial_nearaxis), model='FieldLineAdaptative', initial_conditions=initial_xyz_initial, - maxtime=tmax, times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance) -tracing_optimized = Tracing(field=BiotSavart(coils_optimized), model='FieldLineAdaptative', initial_conditions=initial_xyz_optimized, - maxtime=tmax, times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance) -print(f"Tracing took {time()-time0:.2f} seconds") - -# Plot coils, before and after optimization -fig = plt.figure(figsize=(8, 4)) -ax1 = fig.add_subplot(121, projection='3d') -ax2 = fig.add_subplot(122, projection='3d') -coils_optimized_initial_nearaxis.plot(ax=ax1, show=False) -field_nearaxis_initial.plot(ax=ax1, show=False, alpha=0.35) -tracing_initial.plot(ax=ax1, show=False) -coils_optimized.plot(ax=ax2, show=False) -field_nearaxis_optimized.plot(ax=ax2, show=False, alpha=0.35) -tracing_optimized.plot(ax=ax2, show=False) -plt.show() - -# # Save the coils to a json file -# coils_optimized.to_json("stellarator_coils.json") -# # Load the coils from a json file -# from essos.coils import Coils_from_json -# coils = Coils_from_json("stellarator_coils.json") - -# # Save results in vtk format to analyze in Paraview -# coils_optimized_initial_nearaxsis.to_vtk('coils_initial') -# coils_optimized.to_vtk('coils_optimized') -# tracing_initial.to_vtk('trajectories_initial') -# tracing_optimized.to_vtk('trajectories_final') \ No newline at end of file diff --git a/examples/optimize_coils_and_surface.py b/examples/optimize_coils_and_surface.py deleted file mode 100644 index 1bef5b83..00000000 --- a/examples/optimize_coils_and_surface.py +++ /dev/null @@ -1,254 +0,0 @@ -import os -number_of_processors_to_use = 1 # Parallelization, this should divide ntheta*nphi -os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' -from essos.fields import BiotSavart, near_axis -from essos.dynamics import Particles, Tracing -from essos.surfaces import BdotN_over_B, SurfaceRZFourier, B_on_surface -from essos.coils import Coils, CreateEquallySpacedCurves, Curves -from essos.optimization import optimize_loss_function, new_nearaxis_from_x_and_old_nearaxis -from essos.objective_functions import (loss_coil_curvature, difference_B_gradB_onaxis, - loss_coil_length,field_from_dofs) -import jax.numpy as jnp -from functools import partial -from jax import jit, vmap, devices, device_put, grad, debug -from jax.sharding import Mesh, NamedSharding, PartitionSpec -from time import time -import matplotlib.pyplot as plt - -mesh = Mesh(devices(), ("dev",)) -sharding = NamedSharding(mesh, PartitionSpec("dev", None)) - -ntheta=30 -nphi=30 -input = os.path.join('input_files','input.rotating_ellipse') -surface_initial = SurfaceRZFourier(input, ntheta=ntheta, nphi=nphi, range_torus='half period') - -# Optimization parameters -max_coil_length = 38 -max_coil_curvature = 0.3 -order_Fourier_series_coils = 5 -number_coil_points = order_Fourier_series_coils*10 -maximum_function_evaluations = 20#600 -number_coils_per_half_field_period = 4 -tolerance_optimization = 1e-7 -target_B_on_axis = 5.7 - -nparticles = number_of_processors_to_use -maxtime_tracing = 4e-5 -num_steps=300 -trace_tolerance=1e-5 -model = 'GuidingCenterAdaptative' - -# Initialize coils -current_on_each_coil = 1.714e7 -number_of_field_periods = surface_initial.nfp -major_radius_coils = surface_initial.dofs[0] -minor_radius_coils = major_radius_coils/1.3 -curves = CreateEquallySpacedCurves(n_curves=number_coils_per_half_field_period, - order=order_Fourier_series_coils, - R=major_radius_coils, r=minor_radius_coils, - n_segments=number_coil_points, - nfp=number_of_field_periods, stellsym=True) -coils_initial = Coils(curves=curves, currents=[current_on_each_coil]*number_coils_per_half_field_period) - -# Initialize near-axis -rc=jnp.array([1, 0.045])*major_radius_coils -zs=jnp.array([0,-0.045])*major_radius_coils -etabar=-0.9/major_radius_coils -field_nearaxis_initial = near_axis(rc=rc, zs=zs, etabar=etabar, nfp=number_of_field_periods, B0=target_B_on_axis) - -print(f"Mean Magnetic field on surface: {jnp.mean(jnp.linalg.norm(B_on_surface(surface_initial, BiotSavart(coils_initial)), axis=2))}") - -# Initialize particles -# Xaxis = field_nearaxis_initial.R0*jnp.cos(field_nearaxis_initial.phi) -# Yaxis = field_nearaxis_initial.R0*jnp.sin(field_nearaxis_initial.phi) -# initial_xyz = jnp.array([Xaxis, Yaxis, field_nearaxis_initial.Z0]).T[:nparticles] -# particles = Particles(initial_xyz=initial_xyz, field=BiotSavart(coils_initial)) -# tracing_initial = Tracing(field=coils_initial, particles=particles, maxtime=maxtime_tracing, model=model, timesteps=num_steps) - -# # Plot initial state -# fig = plt.figure(figsize=(9, 8)) -# ax = fig.add_subplot(111, projection='3d') -# tracing_initial.plot(ax=ax, show=False) -# field_nearaxis_initial.plot(r=major_radius_coils/12, ax=ax, show=False) -# coils_initial.plot(ax=ax, show=False) -# surface_initial.plot(ax=ax, show=False) -# plt.show() - -# @partial(jit, static_argnames=['surface','field']) -def grad_AbsB_on_surface(surface, field): - ntheta = surface.ntheta - nphi = surface.nphi - gamma = surface.gamma - gamma_reshaped = gamma.reshape(nphi * ntheta, 3) - gamma_sharded = device_put(gamma_reshaped, sharding) - dAbsB_by_dX_on_surface = jit(vmap(field.dAbsB_by_dX), in_shardings=sharding, out_shardings=sharding)(gamma_sharded) - dAbsB_by_dX_on_surface = dAbsB_by_dX_on_surface.reshape(nphi, ntheta, 3) - return dAbsB_by_dX_on_surface - -# @partial(jit, static_argnames=['field']) -def B_dot_GradAbsB(points, field): - B = field.B(points) - GradAbsB = field.dAbsB_by_dX(points) - B_dot_GradAbsB = jnp.sum(B * GradAbsB, axis=-1) - return B_dot_GradAbsB - -# @partial(jit, static_argnames=['field']) -def grad_B_dot_GradAbsB(points, field): - return grad(B_dot_GradAbsB, argnums=0)(points, field) - -# @partial(jit, static_argnames=['surface','field']) -def grad_B_dot_GradAbsB_on_surface(surface, field): - ntheta = surface.ntheta - nphi = surface.nphi - gamma = surface.gamma - gamma_reshaped = gamma.reshape(nphi * ntheta, 3) - gamma_sharded = device_put(gamma_reshaped, sharding) - partial_grad_B_dot_GradAbsB = partial(grad_B_dot_GradAbsB, field=field) - grad_B_dot_GradAbsB_on_surface = jit(vmap(partial_grad_B_dot_GradAbsB), in_shardings=sharding, out_shardings=sharding)(gamma_sharded) - grad_B_dot_GradAbsB_on_surface = grad_B_dot_GradAbsB_on_surface.reshape(nphi, ntheta, 3) - return grad_B_dot_GradAbsB_on_surface - -# @partial(jit, static_argnames=['surface','field']) -def loss_normal_cross_GradB_dot_grad_B_dot_GradB_surface(surface, field): - gradAbsB_surface = grad_AbsB_on_surface(surface, field) - grad_B_dot_GradB_surface = grad_B_dot_GradAbsB_on_surface(surface, field) - normal_cross_GradB_surface = jnp.cross(surface.normal, gradAbsB_surface, axisa=-1, axisb=-1) - normal_cross_GradB_dot_grad_B_dot_GradB_surface = jnp.sum(normal_cross_GradB_surface * grad_B_dot_GradB_surface, axis=-1) - return normal_cross_GradB_dot_grad_B_dot_GradB_surface - -@partial(jit, static_argnums=(1, 5, 6, 7, 8, 9, 10)) -def loss_coils_and_surface(x, surface_all, field_nearaxis, dofs_curves, currents_scale, nfp, max_coil_length=42, - n_segments=60, stellsym=True, max_coil_curvature=0.5, target_B_on_surface=5.7): - - field=field_from_dofs(x[:-len(surface_all.x)-len(field_nearaxis.x)] ,dofs_curves=dofs_curves, currents_scale=currents_scale, nfp=nfp,n_segments=n_segments, stellsym=stellsym) - surface = SurfaceRZFourier(rc=surface_all.rc, zs=surface_all.zs, nfp=nfp, range_torus=surface_all.range_torus, nphi=surface_all.nphi, ntheta=surface_all.ntheta) - surface.dofs = x[-len(surface_all.x)-len(field_nearaxis.x):-len(field_nearaxis.x)] - - field_nearaxis = new_nearaxis_from_x_and_old_nearaxis(x[-len(field_nearaxis.x):], field_nearaxis) - - coil_length = loss_coil_length(x[:-len(surface_all.x)-len(field_nearaxis.x)],dofs_curves=dofs_curves,currents_scale=currents_scale,nfp=nfp,n_segments=n_segments,stellsym=stellsym,max_coil_length=max_coil_length) - coil_curvature = loss_coil_curvature(x[:-len(surface_all.x)-len(field_nearaxis.x)],dofs_curves=dofs_curves,currents_scale=currents_scale,nfp=nfp,n_segments=n_segments,stellsym=stellsym,max_coil_curvature=max_coil_curvature) - - - coil_length_loss = 1e3*jnp.max(jnp.concatenate([coil_length-max_coil_length,jnp.array([0])])) - coil_curvature_loss = 1e3*jnp.max(jnp.concatenate([coil_curvature-max_coil_curvature,jnp.array([0])])) - - normal_cross_GradB_dot_grad_B_dot_GradB_surface = jnp.sum(jnp.abs(loss_normal_cross_GradB_dot_grad_B_dot_GradB_surface(surface, field))) - - bdotn_over_b = BdotN_over_B(surface, field) - bdotn_over_b_loss = 10*jnp.sum(jnp.abs(bdotn_over_b)) - - mean_cross_sectional_area_loss = 100*jnp.abs(surface.mean_cross_sectional_area()-surface_all.mean_cross_sectional_area()) - - AbsB_on_surface = jnp.linalg.norm(B_on_surface(surface, field), axis=2) - AbsB_surface_loss = jnp.abs(jnp.mean(AbsB_on_surface)-target_B_on_surface) - - B_difference, gradB_difference = difference_B_gradB_onaxis(field_nearaxis, field) - B_difference_loss = 30*jnp.sum(jnp.abs(B_difference)) - gradB_difference_loss = 30*jnp.sum(jnp.abs(gradB_difference)) - - elongation = field_nearaxis.elongation - iota = field_nearaxis.iota - elongation_loss = jnp.sum(jnp.abs(elongation)) - iota_loss = 50/jnp.abs(iota) - - axis_surface = surface.dofs[0] - axis_nearaxis = field_nearaxis.rc[0] - axis_loss = jnp.abs(jnp.min(jnp.array([axis_nearaxis-axis_surface,0]))) - - # debug.print("######################") - # debug.print("normal_cross_GradB_dot_grad_B_dot_GradB_surface={}", normal_cross_GradB_dot_grad_B_dot_GradB_surface) - # debug.print("bdotn_over_b_loss={}", bdotn_over_b_loss) - # debug.print("mean_cross_sectional_area_loss={}", mean_cross_sectional_area_loss) - # debug.print("B_difference_loss={}", B_difference_loss) - # debug.print("gradB_difference_loss={}", gradB_difference_loss) - # debug.print("iota_loss={}", iota_loss) - # debug.print("axis_loss={}", axis_loss) - - # Xaxis = field_nearaxis.R0*jnp.cos(field_nearaxis.phi) - # Yaxis = field_nearaxis.R0*jnp.sin(field_nearaxis.phi) - # initial_xyz = jnp.array([Xaxis, Yaxis, field_nearaxis.Z0]).T[:nparticles] - # particles = Particles(initial_xyz=initial_xyz, field=field) - # particles_drift_loss = jnp.sum(loss_particle_drift(field, particles, maxtime_tracing, num_steps, trace_tolerance, model=model))/num_steps/nparticles - - return ( - coil_length_loss+coil_curvature_loss - # +normal_cross_GradB_dot_grad_B_dot_GradB_surface - +bdotn_over_b_loss - +mean_cross_sectional_area_loss - # +AbsB_surface_loss - +B_difference_loss - +gradB_difference_loss - +elongation_loss - +iota_loss - # +axis_loss - # +particles_drift_loss - ) - -# Optimize coils -print(f'Optimizing coils with {maximum_function_evaluations} function evaluations.') -time0 = time() -initial_dofs = jnp.concatenate((coils_initial.x, surface_initial.x, field_nearaxis_initial.x)) -coils_optimized, surface_optimized, field_nearaxis_optimized = optimize_loss_function(loss_coils_and_surface, initial_dofs=initial_dofs, coils=coils_initial, tolerance_optimization=tolerance_optimization, - maximum_function_evaluations=maximum_function_evaluations, surface_all=surface_initial, field_nearaxis=field_nearaxis_initial, - max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature, target_B_on_surface=target_B_on_axis) -print(f"Optimization took {time()-time0:.2f} seconds") -# Xaxis = field_nearaxis_optimized.R0*jnp.cos(field_nearaxis_optimized.phi) -# Yaxis = field_nearaxis_optimized.R0*jnp.sin(field_nearaxis_optimized.phi) -# initial_xyz = jnp.array([Xaxis, Yaxis, field_nearaxis_optimized.Z0]).T[:nparticles] -# particles = Particles(initial_xyz=initial_xyz, field=BiotSavart(coils_optimized)) -# tracing_optimized = Tracing(field=coils_optimized, particles=particles, maxtime=maxtime_tracing, model=model, timesteps=num_steps) - -print(f'############################################') -print(f"Mean Magnetic field on surface: {jnp.mean(jnp.linalg.norm(B_on_surface(surface_optimized, BiotSavart(coils_optimized)), axis=2))}") -print(f"Initial max(BdotN/B): {jnp.max(BdotN_over_B(surface_initial, BiotSavart(coils_initial))):.2e}") -print(f"Optimized max(BdotN/B): {jnp.max(BdotN_over_B(surface_optimized, BiotSavart(coils_optimized))):.2e}") -print(f'Initial iota on-axis: {field_nearaxis_initial.iota}') -print(f'Optimized iota on-axis: {field_nearaxis_optimized.iota}') -print(f'Initial max(elongation): {max(field_nearaxis_initial.elongation)}') -print(f'Optimized max(elongation): {max(field_nearaxis_optimized.elongation)}') -print(f"Initial coils length: {coils_initial.length[:number_coils_per_half_field_period]}") -print(f"Optimized coils length: {coils_optimized.length[:number_coils_per_half_field_period]}") -print(f"Initial coils curvature: {jnp.mean(coils_initial.curvature, axis=1)[:number_coils_per_half_field_period]}") -print(f"Optimized coils curvature: {jnp.mean(coils_optimized.curvature, axis=1)[:number_coils_per_half_field_period]}") -B_difference_initial, gradB_difference_initial = difference_B_gradB_onaxis(field_nearaxis_initial, BiotSavart(coils_initial)) -B_difference_optimized, gradB_difference_optimized = difference_B_gradB_onaxis(field_nearaxis_optimized, BiotSavart(coils_optimized)) -print(f'Initial B on axis difference: {jnp.sum(jnp.abs(B_difference_initial))}') -print(f'Optimized B on axis difference: {jnp.sum(jnp.abs(B_difference_optimized))}') -print(f'Initial gradB on axis difference: {jnp.sum(jnp.abs(gradB_difference_initial))}') -print(f'Optimized gradB on axis difference: {jnp.sum(jnp.abs(gradB_difference_optimized))}') - -# Plot coils, before and after optimization -fig = plt.figure(figsize=(8, 4)) -ax1 = fig.add_subplot(121, projection='3d') -ax2 = fig.add_subplot(122, projection='3d') -coils_initial.plot(ax=ax1, show=False) -surface_initial.plot(ax=ax1, show=False) -field_nearaxis_initial.plot(r=major_radius_coils/12, ax=ax1, show=False) -# tracing_initial.plot(ax=ax1, show=False) -coils_optimized.plot(ax=ax2, show=False) -surface_optimized.plot(ax=ax2, show=False) -field_nearaxis_optimized.plot(r=major_radius_coils/12, ax=ax2, show=False) -# tracing_optimized.plot(ax=ax2, show=False) -plt.tight_layout() -plt.show() - -# Save the surface to a VMEC file -surface_optimized.to_vmec('input.optimized') - -# Save results in vtk format to analyze in Paraview -surface_initial.to_vtk('initial_surface', field=BiotSavart(coils_initial)) -coils_initial.to_vtk('initial_coils') -field_nearaxis_initial.to_vtk('initial_field_nearaxis', r=major_radius_coils/12, field=BiotSavart(coils_initial)) -surface_optimized.to_vtk('optimized_surface', field=BiotSavart(coils_optimized)) -coils_optimized.to_vtk('optimized_coils') -field_nearaxis_optimized.to_vtk('optimized_field_nearaxis', r=major_radius_coils/12, field=BiotSavart(coils_optimized)) -# tracing_initial.to_vtk('initial_tracing') -# tracing_optimized.to_vtk('optimized_tracing') - -# # Save the coils to a json file -# coils_optimized.to_json("stellarator_coils.json") -# # Load the coils from a json file -# from essos.coils import Coils_from_json -# coils = Coils_from_json("stellarator_coils.json") \ No newline at end of file diff --git a/examples/optimize_coils_for_nearaxis.py b/examples/optimize_coils_for_nearaxis.py deleted file mode 100644 index 6d70ea04..00000000 --- a/examples/optimize_coils_for_nearaxis.py +++ /dev/null @@ -1,88 +0,0 @@ -from time import time -import jax.numpy as jnp -import matplotlib.pyplot as plt -from essos.coils import Coils, CreateEquallySpacedCurves -from essos.fields import near_axis, BiotSavart -from essos.dynamics import Tracing -from essos.optimization import optimize_loss_function -from essos.objective_functions import loss_coils_for_nearaxis - -# Optimization parameters -max_coil_length = 4 -max_coil_curvature = 6 -order_Fourier_series_coils = 5 -number_coil_points = order_Fourier_series_coils*10 -maximum_function_evaluations = 200 -number_coils_per_half_field_period = 3 -tolerance_optimization = 1e-8 - -# Initialize Near-Axis field -rc=jnp.array([1, 0.045]) -zs=jnp.array([0,-0.045]) -etabar=-0.9 -nfp=3 -field = near_axis(rc=rc, zs=zs, etabar=etabar, nfp=nfp) - -# Initialize coils -current_on_each_coil = 17e5*field.B0/nfp/2 -number_of_field_periods = nfp -major_radius_coils = field.R0[0] -minor_radius_coils = major_radius_coils/2.0 -curves = CreateEquallySpacedCurves(n_curves=number_coils_per_half_field_period, - order=order_Fourier_series_coils, - R=major_radius_coils, r=minor_radius_coils, - n_segments=number_coil_points, - nfp=number_of_field_periods, stellsym=True) -coils_initial = Coils(curves=curves, currents=[current_on_each_coil]*number_coils_per_half_field_period) - -# Optimize coils -print(f'Optimizing coils with {maximum_function_evaluations} function evaluations.') -time0 = time() -initial_dofs = coils_initial.x -coils_optimized = optimize_loss_function(loss_coils_for_nearaxis, initial_dofs=coils_initial.x, - coils=coils_initial, tolerance_optimization=tolerance_optimization, - maximum_function_evaluations=maximum_function_evaluations, field_nearaxis=field, - max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature,) -print(f"Optimization took {time()-time0:.2f} seconds") - -# Trace fieldlines -nfieldlines = 6 -num_steps = 1000 -tmax = 1.e-6 -trace_tolerance = 1e-7 - -R0 = jnp.linspace(field.R0[0], 1.05*field.R0[0], nfieldlines) -Z0 = jnp.zeros(nfieldlines) -phi0 = jnp.zeros(nfieldlines) -initial_xyz=jnp.array([R0*jnp.cos(phi0), R0*jnp.sin(phi0), Z0]).T - -time0 = time() -tracing_initial = Tracing(field=BiotSavart(coils_initial), model='FieldLineAdaptative', initial_conditions=initial_xyz, - maxtime=tmax, times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance) -tracing_optimized = Tracing(field=BiotSavart(coils_optimized), model='FieldLineAdaptative', initial_conditions=initial_xyz, - maxtime=tmax, times_to_trace=num_steps, atol=trace_tolerance,rtol=trace_tolerance) -print(f"Tracing took {time()-time0:.2f} seconds") - -# Plot coils, before and after optimization -fig = plt.figure(figsize=(8, 4)) -ax1 = fig.add_subplot(121, projection='3d') -ax2 = fig.add_subplot(122, projection='3d') -coils_initial.plot(ax=ax1, show=False) -field.plot(ax=ax1, show=False, alpha=0.2) -tracing_initial.plot(ax=ax1, show=False) -coils_optimized.plot(ax=ax2, show=False) -field.plot(ax=ax2, show=False, alpha=0.2) -tracing_optimized.plot(ax=ax2, show=False) -plt.show() - -# # Save the coils to a json file -# coils_optimized.to_json("stellarator_coils.json") -# # Load the coils from a json file -# from essos.coils import Coils_from_json -# coils = Coils_from_json("stellarator_coils.json") - -# # Save results in vtk format to analyze in Paraview -# coils_initial.to_vtk('coils_initial') -# coils_optimized.to_vtk('coils_optimized') -# tracing_initial.to_vtk('trajectories_initial') -# tracing_optimized.to_vtk('trajectories_final') \ No newline at end of file diff --git a/examples/optimize_coils_particle_confinement_fullorbit.py b/examples/optimize_coils_particle_confinement_fullorbit.py deleted file mode 100644 index e49fe27f..00000000 --- a/examples/optimize_coils_particle_confinement_fullorbit.py +++ /dev/null @@ -1,99 +0,0 @@ - -import os -number_of_processors_to_use = 1 # Parallelization, this should divide nparticles -os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' -from time import time -import jax.numpy as jnp -import matplotlib.pyplot as plt -from essos.dynamics import Particles, Tracing -from essos.fields import BiotSavart -from essos.coils import Coils, CreateEquallySpacedCurves -from essos.optimization import optimize_loss_function -from essos.objective_functions import loss_optimize_coils_for_particle_confinement - -# Optimization parameters -target_B_on_axis = 5.7 -max_coil_length = 31 -max_coil_curvature = 0.4 -nparticles = number_of_processors_to_use*1 -order_Fourier_series_coils = 4 -number_coil_points = 80 -maximum_function_evaluations = 10 -maxtime_tracing = 1e-6 -number_coils_per_half_field_period = 3 -number_of_field_periods = 2 -model = 'FullOrbit_Boris' -timesteps = 3000#int(3*maxtime_tracing/1e-8) - -nparticles_plot = number_of_processors_to_use*2 -model_plot = 'GuidingCenterAdaptative' -timesteps_plot = 10000 -maxtime_tracing_plot = 3e-5 - -# Initialize coils -current_on_each_coil = 1.84e7 -major_radius_coils = 7.75 -minor_radius_coils = 4.5 -curves = CreateEquallySpacedCurves(n_curves=number_coils_per_half_field_period, - order=order_Fourier_series_coils, - R=major_radius_coils, r=minor_radius_coils, - n_segments=number_coil_points, - nfp=number_of_field_periods, stellsym=True) -coils_initial = Coils(curves=curves, currents=[current_on_each_coil]*number_coils_per_half_field_period) - -# Initialize particles -phi_array = jnp.linspace(0, 2*jnp.pi, nparticles) -initial_xyz=jnp.array([major_radius_coils*jnp.cos(phi_array), major_radius_coils*jnp.sin(phi_array), 0*phi_array]).T -particles = Particles(initial_xyz=initial_xyz) -particles.to_full_orbit(BiotSavart(coils_initial)) -tracing_initial = Tracing(field=coils_initial, particles=particles, maxtime=maxtime_tracing, model=model, times_to_trace=timesteps) - -# Optimize coils -print(f'Optimizing coils with {maximum_function_evaluations} function evaluations and maxtime_tracing={maxtime_tracing}') -time0 = time() -coils_optimized = optimize_loss_function(loss_optimize_coils_for_particle_confinement, initial_dofs=coils_initial.x, - coils=coils_initial, tolerance_optimization=1e-4, particles=particles, - maximum_function_evaluations=maximum_function_evaluations, max_coil_curvature=max_coil_curvature, - target_B_on_axis=target_B_on_axis, max_coil_length=max_coil_length, model=model, - maxtime=maxtime_tracing, num_steps=timesteps) -print(f" Optimization took {time()-time0:.2f} seconds") -particles.to_full_orbit(BiotSavart(coils_optimized)) - -phi_array_plot = jnp.linspace(0, 2*jnp.pi, nparticles_plot) -initial_xyz_plot=jnp.array([major_radius_coils*jnp.cos(phi_array_plot), major_radius_coils*jnp.sin(phi_array_plot), 0*phi_array_plot]).T -particles_plot = Particles(initial_xyz=initial_xyz_plot) -particles.to_full_orbit(BiotSavart(coils_optimized)) -tracing_optimized = Tracing(field=coils_optimized, particles=particles, maxtime=maxtime_tracing_plot, model=model_plot, times_to_trace=timesteps_plot) - -# Plot trajectories, before and after optimization -fig = plt.figure(figsize=(9, 8)) -ax1 = fig.add_subplot(221, projection='3d') -ax2 = fig.add_subplot(222, projection='3d') -ax3 = fig.add_subplot(223) -ax4 = fig.add_subplot(224) - -coils_initial.plot(ax=ax1, show=False) -tracing_initial.plot(ax=ax1, show=False) -for i, trajectory in enumerate(tracing_initial.trajectories): - ax3.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}', linewidth=0.2) -ax3.set_xlabel('R (m)');ax3.set_ylabel('Z (m)');#ax3.legend() -coils_optimized.plot(ax=ax2, show=False) -tracing_optimized.plot(ax=ax2, show=False) -# for i, trajectory in enumerate(tracing_optimized.trajectories): -# ax4.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}', linewidth=0.2) -# ax4.set_xlabel('R (m)');ax4.set_ylabel('Z (m)');#ax4.legend() -plotting_data = tracing_optimized.poincare_plot(ax=ax4, shifts = [jnp.pi/4, jnp.pi/2, 3*jnp.pi/4], show=False) -plt.tight_layout() -plt.show() - -# # Save the coils to a json file -# coils_optimized.to_json("stellarator_coils.json") -# # Load the coils from a json file -# from essos.coils import Coils_from_json -# coils = Coils_from_json("stellarator_coils.json") - -# # Save results in vtk format to analyze in Paraview -# tracing_initial.to_vtk('trajectories_initial') -# tracing_optimized.to_vtk('trajectories_final') -# coils_initial.to_vtk('coils_initial') -# coils_optimized.to_vtk('coils_optimized') \ No newline at end of file diff --git a/examples/optimize_coils_particle_confinement_guidingcenter_adam.py b/examples/optimize_coils_particle_confinement_guidingcenter_adam.py deleted file mode 100644 index f5e66ab3..00000000 --- a/examples/optimize_coils_particle_confinement_guidingcenter_adam.py +++ /dev/null @@ -1,125 +0,0 @@ - -import os -number_of_processors_to_use = 1 # Parallelization, this should divide nparticles -os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' -import jax.numpy as jnp -from jax import jit, grad -import matplotlib.pyplot as plt -from essos.dynamics import Particles, Tracing -from essos.coils import Coils, CreateEquallySpacedCurves,Curves -from essos.objective_functions import loss_particle_r_cross_max_constraint -from essos.objective_functions import loss_coil_curvature,loss_coil_length, loss_normB_axis_average -from functools import partial -import optax - - -# Optimization parameters -target_B_on_axis = 5.7 -max_coil_length = 31 -max_coil_curvature = 0.4 -nparticles = number_of_processors_to_use*1 -order_Fourier_series_coils = 4 -number_coil_points = 80 -maximum_function_evaluations = 15 -maxtimes = [2.e-5] -num_steps=100 -number_coils_per_half_field_period = 3 -number_of_field_periods = 2 -model = 'GuidingCenterAdaptative' - -# Initialize coils -current_on_each_coil = 1.84e7 -major_radius_coils = 7.75 -minor_radius_coils = 4.45 -curves = CreateEquallySpacedCurves(n_curves=number_coils_per_half_field_period, - order=order_Fourier_series_coils, - R=major_radius_coils, r=minor_radius_coils, - n_segments=number_coil_points, - nfp=number_of_field_periods, stellsym=True) -coils_initial = Coils(curves=curves, currents=[current_on_each_coil]*number_coils_per_half_field_period) - -len_dofs_curves = len(jnp.ravel(coils_initial.dofs_curves)) -nfp = coils_initial.nfp -stellsym = coils_initial.stellsym -n_segments = coils_initial.n_segments -dofs_curves_shape = coils_initial.dofs_curves.shape -currents_scale = coils_initial.currents_scale - -# Initialize particles -phi_array = jnp.linspace(0, 2*jnp.pi, nparticles) -initial_xyz=jnp.array([major_radius_coils*jnp.cos(phi_array), major_radius_coils*jnp.sin(phi_array), 0*phi_array]).T -particles = Particles(initial_xyz=initial_xyz) - -t=maxtimes[0] - -curvature_partial=partial(loss_coil_curvature, dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,max_coil_curvature=max_coil_curvature) -length_partial=partial(loss_coil_length, dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,max_coil_length=max_coil_length) -Baxis_average_partial=partial(loss_normB_axis_average,dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,npoints=15,target_B_on_axis=target_B_on_axis) -r_max_partial = partial(loss_particle_r_cross_max_constraint, particles=particles,dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,maxtime=t,model = model,num_steps=num_steps) - -params=coils_initial.x -optimizer=optax.adabelief(learning_rate=0.003) -opt_state=optimizer.init(params) - -def total_loss(params): - return jnp.linalg.norm(curvature_partial(params)+length_partial(params)+Baxis_average_partial(params))**2 - - -jit -def update(params,opt_state): - this_grad = grad(total_loss)(params) - updates, opt_state =optimizer.update(this_grad, opt_state) - params = optax.apply_updates(params, updates) - return params,opt_state - -for i in range(maximum_function_evaluations): - params,opt_state=update(params,opt_state) - if i % 3 == 0: - print('Objective function at iteration {:d}: {:.2E}'.format(i, total_loss(params))) - - -dofs_curves = jnp.reshape(params[:len_dofs_curves], (dofs_curves_shape)) -dofs_currents = params[len_dofs_curves:] -curves = Curves(dofs_curves, n_segments, nfp, stellsym) -new_coils = Coils(curves=curves, currents=dofs_currents*coils_initial.currents_scale) -params=new_coils.x -tracing_initial = Tracing(field=coils_initial, particles=particles, maxtime=t, model=model - ,times_to_trace=200,timestep=1.e-8,atol=1.e-5,rtol=1.e-5) -tracing_optimized = Tracing(field=new_coils, particles=particles, maxtime=t, model=model,times_to_trace=200,timestep=1.e-8,atol=1.e-5,rtol=1.e-5) - - -# Plot trajectories, before and after optimization -fig = plt.figure(figsize=(9, 8)) -ax1 = fig.add_subplot(221, projection='3d') -ax2 = fig.add_subplot(222, projection='3d') -ax3 = fig.add_subplot(223) -ax4 = fig.add_subplot(224) - -coils_initial.plot(ax=ax1, show=False) -tracing_initial.plot(ax=ax1, show=False) -for i, trajectory in enumerate(tracing_initial.trajectories): - ax3.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') - -ax3.set_xlabel('R (m)') -ax3.set_ylabel('Z (m)') -#ax3.legend() -new_coils.plot(ax=ax2, show=False) -tracing_optimized.plot(ax=ax2, show=False) -for i, trajectory in enumerate(tracing_optimized.trajectories): - ax4.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') -ax4.set_xlabel('R (m)') -ax4.set_ylabel('Z (m)')#ax4.legend() -plt.tight_layout() -# plt.savefig(f'opt_adam.pdf') -plt.savefig('optimize_coils_particle_confinement_guidingcenter_adam.png', dpi=300) -# # Save the coils to a json file -# coils_optimized.to_json("stellarator_coils.json") -# # Load the coils from a json file -# from essos.coils import Coils_from_json -# coils = Coils_from_json("stellarator_coils.json") - -# # Save results in vtk format to analyze in Paraview -# tracing_initial.to_vtk('trajectories_initial') -#tracing_optimized.to_vtk('trajectories_final') -#coils_initial.to_vtk('coils_initial') -#new_coils.to_vtk('coils_optimized') \ No newline at end of file diff --git a/examples/optimize_coils_particle_confinement_guidingcenter_augmented_lagrangian.py b/examples/optimize_coils_particle_confinement_guidingcenter_augmented_lagrangian.py deleted file mode 100644 index d764f8bf..00000000 --- a/examples/optimize_coils_particle_confinement_guidingcenter_augmented_lagrangian.py +++ /dev/null @@ -1,174 +0,0 @@ - -import os -number_of_processors_to_use = 1 # Parallelization, this should divide nparticles -os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' -from time import time -import jax -print(jax.devices()) -jax.config.update("jax_enable_x64", True) -import jax.numpy as jnp -import matplotlib.pyplot as plt -from essos.surfaces import SurfaceRZFourier -from essos.dynamics import Particles, Tracing -from essos.coils import Coils, CreateEquallySpacedCurves,Curves -from essos.objective_functions import loss_particle_r_cross_max_constraint,loss_particle_gamma_c -from essos.objective_functions import loss_coil_curvature,loss_coil_length,loss_normB_axis_average,loss_Br,loss_iota -from functools import partial -import essos.augmented_lagrangian as alm - - - - - -# Optimization parameters -target_B_on_axis = 5.7 -max_coil_length = 31 -max_coil_curvature = 0.4 -nparticles = number_of_processors_to_use*10 -order_Fourier_series_coils = 4 -number_coil_points = 80 -maximum_function_evaluations = 1 -maxtimes = [1.e-5] -num_steps=100 -number_coils_per_half_field_period = 3 -number_of_field_periods = 2 -model = 'GuidingCenter' - -# Initialize coils -current_on_each_coil = 1.84e7 -major_radius_coils = 7.75 -minor_radius_coils = 4.45 -curves = CreateEquallySpacedCurves(n_curves=number_coils_per_half_field_period, - order=order_Fourier_series_coils, - R=major_radius_coils, r=minor_radius_coils, - n_segments=number_coil_points, - nfp=number_of_field_periods, stellsym=True) -coils_initial = Coils(curves=curves, currents=[current_on_each_coil]*number_coils_per_half_field_period) - - -len_dofs_curves = len(jnp.ravel(coils_initial.dofs_curves)) -nfp = coils_initial.nfp -stellsym = coils_initial.stellsym -n_segments = coils_initial.n_segments -dofs_curves_shape = coils_initial.dofs_curves.shape -currents_scale = coils_initial.currents_scale - - -# Initialize particles -phi_array = jnp.linspace(0, 2*jnp.pi, nparticles) -initial_xyz=jnp.array([major_radius_coils*jnp.cos(phi_array), major_radius_coils*jnp.sin(phi_array), 0*phi_array]).T -particles = Particles(initial_xyz=initial_xyz) - -t=maxtimes[0] -loss_partial = partial(loss_particle_gamma_c,particles=particles, dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,maxtime=t,model=model,num_steps=num_steps) -curvature_partial=partial(loss_coil_curvature, dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,max_coil_curvature=max_coil_curvature) -length_partial=partial(loss_coil_length, dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,max_coil_length=max_coil_length) -Baxis_average_partial=partial(loss_normB_axis_average,dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,npoints=15,target_B_on_axis=target_B_on_axis) -r_max_partial = partial(loss_particle_r_cross_max_constraint,target_r=0.4, particles=particles,dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,maxtime=t,model=model,num_steps=num_steps) -iota_partial = partial(loss_iota,target_iota=0.5, particles=particles,dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,maxtime=t,model=model,num_steps=num_steps) -Br_partial = partial(loss_Br, particles=particles,dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,maxtime=t,model=model,num_steps=num_steps) - - -# Create the constraints -penalty = 1. #Intial penalty values -multiplier=0.5 #Initial lagrange multiplier values -sq_grad=0.0 #Initial square gradient parameter value for Mu adaptative -model_lagrangian='Standard' #Use standard augmented lagragian suitable for bounded optimizers -#Since we are using LBFGS-B from jaxopt, model_mu will be updated with tolerances so we do not need to difinte the model - -#Construct constraints -constraints = alm.combine( -alm.eq(curvature_partial, multiplier=multiplier,penalty=penalty,sq_grad=sq_grad), -alm.eq(length_partial, multiplier=multiplier,penalty=penalty,sq_grad=sq_grad), -alm.eq(Baxis_average_partial, multiplier=multiplier,penalty=penalty,sq_grad=sq_grad), -alm.eq(r_max_partial, multiplier=multiplier,penalty=penalty,sq_grad=sq_grad), -#alm.eq(Br_partial, multiplier=multiplier,penalty=penalty,sq_grad=sq_grad), -#alm.eq(iota_partial, multiplier=multiplier,penalty=penalty,sq_grad=sq_grad), -) - - - -beta=2. #penalty update parameter -mu_max=1.e4 #Maximum penalty parameter allowed -alpha=0.99 #These are parameters only used if gradient descent and adaaptative mu -gamma=1.e-2 -epsilon=1.e-8 -omega_tol=0.0001 #desired grad_tolerance, associated with grad of lagrangian to main parameters -eta_tol=0.001 #desired contraint tolerance, associated with variation of contraints - - - -#If loss=cost_function(x) is not prescribed, f(x)=0 is considered -ALM=alm.ALM_model_jaxopt_lbfgsb(constraints,loss=loss_partial,model_lagrangian=model_lagrangian,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) - -#Initializing lagrange multipliers -lagrange_params=constraints.init(coils_initial.x) -#parameters are a tuple of the primal/main optimisation parameters and the lagrange multipliers -params = coils_initial.x, lagrange_params -#This is just to initialize an empty state for the lagrange multiplier update and get some information -lag_state,grad,info=ALM.init(params) - -#Initializing first tolerances for the inner minimisation loop iteration -mu_average=alm.penalty_average(lagrange_params) -#omega=1.#1./mu_average -#eta=1000.#1./mu_average**0.1 -omega=1./mu_average -eta=1./mu_average**0.1 - -i=0 -while i<=maximum_function_evaluations and (jnp.linalg.norm(grad[0])>omega_tol or alm.norm_constraints(info[2])>eta_tol): - #One step of ALM optimization - params, lag_state,grad,info,eta,omega = ALM.update(params,lag_state,grad,info,eta,omega) - #if i % 5 == 0: - #print(f'i: {i}, loss f: {info[0]:g}, infeasibility: {alm.total_infeasibility(info[1]):g}') - print(f'i: {i}, loss f: {info[0]:g},loss L: {info[1]:g}, infeasibility: {alm.total_infeasibility(info[2]):g}') - #print('lagrange',params[1]) - i=i+1 - - -dofs_curves = jnp.reshape(params[0][:len_dofs_curves], (dofs_curves_shape)) -dofs_currents = params[0][len_dofs_curves:] -curves = Curves(dofs_curves, n_segments, nfp, stellsym) -new_coils = Coils(curves=curves, currents=dofs_currents*coils_initial.currents_scale) -params=new_coils.x -tracing_initial = Tracing(field=coils_initial, particles=particles, maxtime=t, model=model - ,times_to_trace=200,timestep=1.e-8,atol=1.e-5,rtol=1.e-5) -tracing_optimized = Tracing(field=new_coils, particles=particles, maxtime=t, model=model,times_to_trace=200,timestep=1.e-8,atol=1.e-5,rtol=1.e-5) - -#print('Final params',params) -#print(info[1]) -# Plot trajectories, before and after optimization -fig = plt.figure(figsize=(9, 8)) -ax1 = fig.add_subplot(221, projection='3d') -ax2 = fig.add_subplot(222, projection='3d') -ax3 = fig.add_subplot(223) -ax4 = fig.add_subplot(224) - -coils_initial.plot(ax=ax1, show=False) -tracing_initial.plot(ax=ax1, show=False) -for i, trajectory in enumerate(tracing_initial.trajectories): - ax3.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') - -ax3.set_xlabel('R (m)') -ax3.set_ylabel('Z (m)') -#ax3.legend() -new_coils.plot(ax=ax2, show=False) -tracing_optimized.plot(ax=ax2, show=False) -for i, trajectory in enumerate(tracing_optimized.trajectories): - ax4.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') -ax4.set_xlabel('R (m)') -ax4.set_ylabel('Z (m)')#ax4.legend() -plt.tight_layout() -plt.savefig(f'opt_constrained.pdf') - -# # Save the coils to a json file -# coils_optimized.to_json("stellarator_coils.json") -# # Load the coils from a json file -# from essos.coils import Coils_from_json -# coils = Coils_from_json("stellarator_coils.json") - -# # Save results in vtk format to analyze in Paraview -# tracing_initial.to_vtk('trajectories_initial') -#tracing_optimized.to_vtk('trajectories_final') -#coils_initial.to_vtk('coils_initial') -#new_coils.to_vtk('coils_optimized') diff --git a/examples/optimize_coils_particle_confinement_guidingcenter_lbfgs.py b/examples/optimize_coils_particle_confinement_guidingcenter_lbfgs.py deleted file mode 100644 index 2ee77702..00000000 --- a/examples/optimize_coils_particle_confinement_guidingcenter_lbfgs.py +++ /dev/null @@ -1,126 +0,0 @@ - -import os -number_of_processors_to_use = 1 # Parallelization, this should divide nparticles -os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' -from jax import jit, value_and_grad -import jax.numpy as jnp -import matplotlib.pyplot as plt -from essos.dynamics import Particles, Tracing -from essos.coils import Coils, CreateEquallySpacedCurves,Curves -from essos.optimization import optimize_loss_function -from essos.objective_functions import loss_particle_r_cross_final,loss_particle_radial_drift,loss_particle_gamma_c -from essos.objective_functions import loss_coil_curvature_new,loss_coil_length_new,loss_normB_axis,loss_normB_axis_average -from functools import partial -import optax - - -# Optimization parameters -target_B_on_axis = 5.7 -max_coil_length = 31 -max_coil_curvature = 0.4 -nparticles = number_of_processors_to_use*10 -order_Fourier_series_coils = 4 -number_coil_points = 80 -maximum_function_evaluations = 3 -maxtimes = [2.e-5] -num_steps=100 -number_coils_per_half_field_period = 3 -number_of_field_periods = 2 -model = 'GuidingCenterAdaptative' - -# Initialize coils -current_on_each_coil = 1.84e7 -major_radius_coils = 7.75 -minor_radius_coils = 4.45 -curves = CreateEquallySpacedCurves(n_curves=number_coils_per_half_field_period, - order=order_Fourier_series_coils, - R=major_radius_coils, r=minor_radius_coils, - n_segments=number_coil_points, - nfp=number_of_field_periods, stellsym=True) -coils_initial = Coils(curves=curves, currents=[current_on_each_coil]*number_coils_per_half_field_period) - -len_dofs_curves = len(jnp.ravel(coils_initial.dofs_curves)) -nfp = coils_initial.nfp -stellsym = coils_initial.stellsym -n_segments = coils_initial.n_segments -dofs_curves_shape = coils_initial.dofs_curves.shape -currents_scale = coils_initial.currents_scale - -# Initialize particles -phi_array = jnp.linspace(0, 2*jnp.pi, nparticles) -initial_xyz=jnp.array([major_radius_coils*jnp.cos(phi_array), major_radius_coils*jnp.sin(phi_array), 0*phi_array]).T -particles = Particles(initial_xyz=initial_xyz) - -t=maxtimes[0] - -curvature_partial=partial(loss_coil_curvature_new, dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,max_coil_curvature=max_coil_curvature) -length_partial=partial(loss_coil_length_new, dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,max_coil_length=max_coil_length) -Baxis_average_partial=partial(loss_normB_axis_average,dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,npoints=15,target_B_on_axis=target_B_on_axis) -r_max_partial = partial(loss_particle_r_cross_final, particles=particles,dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,maxtime=t,model = model,num_steps=num_steps) -def total_loss(params): - return jnp.linalg.norm(jnp.concatenate((jnp.ravel(r_max_partial(params)),length_partial(params),curvature_partial(params),Baxis_average_partial(params))))**2 - -params=coils_initial.x -optimizer=optax.lbfgs() -opt_state=optimizer.init(params) - -@jit -def update(params,opt_state): - value, grad = value_and_grad(total_loss)(params) - updates, opt_state =optimizer.update(grad, opt_state, params, value=value, grad=grad, value_fn=total_loss) - params = optax.apply_updates(params, updates) - return params,opt_state - -for i in range(maximum_function_evaluations): - params,opt_state=update(params,opt_state) - if i % 3 == 0: - print('Objective function at iteration {:d}: {:.2E}'.format(i, total_loss(params))) - -dofs_curves = jnp.reshape(params[:len_dofs_curves], (dofs_curves_shape)) -dofs_currents = params[len_dofs_curves:] -curves = Curves(dofs_curves, n_segments, nfp, stellsym) -new_coils = Coils(curves=curves, currents=dofs_currents*coils_initial.currents_scale) -params=new_coils.x -tracing_initial = Tracing(field=coils_initial, particles=particles, maxtime=t, model=model - ,times_to_trace=200,timestep=1.e-8,atol=1.e-5,rtol=1.e-5) -tracing_optimized = Tracing(field=new_coils, particles=particles, maxtime=t, model=model,times_to_trace=200,timestep=1.e-8,atol=1.e-5,rtol=1.e-5) - - -# Plot trajectories, before and after optimization -fig = plt.figure(figsize=(9, 8)) -ax1 = fig.add_subplot(221, projection='3d') -ax2 = fig.add_subplot(222, projection='3d') -ax3 = fig.add_subplot(223) -ax4 = fig.add_subplot(224) - -coils_initial.plot(ax=ax1, show=False) -tracing_initial.plot(ax=ax1, show=False) -for i, trajectory in enumerate(tracing_initial.trajectories): - ax3.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') - -ax3.set_xlabel('R (m)') -ax3.set_ylabel('Z (m)') -#ax3.legend() -new_coils.plot(ax=ax2, show=False) -tracing_optimized.plot(ax=ax2, show=False) -for i, trajectory in enumerate(tracing_optimized.trajectories): - ax4.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') -ax4.set_xlabel('R (m)') -ax4.set_ylabel('Z (m)')#ax4.legend() -plt.tight_layout() -# plt.savefig(f'opt_lbfgs.pdf') -plt.show() - -# # Save the coils to a json file -# coils_optimized.to_json("stellarator_coils.json") -# # Load the coils from a json file -# from essos.coils import Coils_from_json -# coils = Coils_from_json("stellarator_coils.json") - -# # Save results in vtk format to analyze in Paraview -# tracing_initial.to_vtk('trajectories_initial') -#tracing_optimized.to_vtk('trajectories_final') -#coils_initial.to_vtk('coils_initial') -#new_coils.to_vtk('coils_optimized') - - diff --git a/examples/optimize_coils_particle_confinement_loss_fraction_augmented_lagrangian.py b/examples/optimize_coils_particle_confinement_loss_fraction_augmented_lagrangian.py deleted file mode 100644 index 6d35c883..00000000 --- a/examples/optimize_coils_particle_confinement_loss_fraction_augmented_lagrangian.py +++ /dev/null @@ -1,180 +0,0 @@ - -import os -number_of_processors_to_use = 8 # Parallelization, this should divide nparticles -os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' -from time import time -import jax -print(jax.devices()) -jax.config.update("jax_enable_x64", True) -import jax.numpy as jnp -import matplotlib.pyplot as plt -from essos.surfaces import SurfaceRZFourier, SurfaceClassifier -from essos.dynamics import Particles, Tracing -from essos.coils import Coils, CreateEquallySpacedCurves,Curves -from essos.objective_functions import loss_lost_fraction,loss_lost_fraction_times -from essos.objective_functions import loss_coil_curvature,loss_coil_length,loss_normB_axis_average,loss_Br,loss_iota -from functools import partial -import essos.augmented_lagrangian as alm - - - - - -# Optimization parameters -target_B_on_axis = 5.7 -max_coil_length = 31 -max_coil_curvature = 0.4 -nparticles = number_of_processors_to_use*1 -order_Fourier_series_coils = 4 -number_coil_points = 80 -maximum_function_evaluations = 10 -maxtimes = [1.e-2] -timestep=1.e-8 -num_steps=100 -number_coils_per_half_field_period = 3 -number_of_field_periods = 2 -model = 'GuidingCenterAdaptative' - -# Initialize coils -current_on_each_coil = 1.84e7 -major_radius_coils = 7.75 -minor_radius_coils = 4.45 -curves = CreateEquallySpacedCurves(n_curves=number_coils_per_half_field_period, - order=order_Fourier_series_coils, - R=major_radius_coils, r=minor_radius_coils, - n_segments=number_coil_points, - nfp=number_of_field_periods, stellsym=True) -coils_initial = Coils(curves=curves, currents=[current_on_each_coil]*number_coils_per_half_field_period) - - -len_dofs_curves = len(jnp.ravel(coils_initial.dofs_curves)) -nfp = coils_initial.nfp -stellsym = coils_initial.stellsym -n_segments = coils_initial.n_segments -dofs_curves_shape = coils_initial.dofs_curves.shape -currents_scale = coils_initial.currents_scale - -ntheta=30 -nphi=30 -input = os.path.join(os.path.dirname(__name__),'input_files','input.toroidal_surface') -surface= SurfaceRZFourier(input, ntheta=ntheta, nphi=nphi, range_torus='full torus') -timeI=time() -boundary=SurfaceClassifier(surface,h=0.1) -print(f"ESSOS boundary took {time()-timeI:.2f} seconds") -#print('Final params',params) -#print(info[1]) -# Plot trajectories, before and after optimization - - -# Initialize particles -phi_array = jnp.linspace(0, 2*jnp.pi, nparticles) -initial_xyz=jnp.array([major_radius_coils*jnp.cos(phi_array), major_radius_coils*jnp.sin(phi_array), 0*phi_array]).T -particles = Particles(initial_xyz=initial_xyz) - -t=maxtimes[0] -loss_partial = partial(loss_lost_fraction,particles=particles, dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,maxtime=t,timestep=timestep,model=model,num_steps=num_steps,boundary=boundary) -jax.grad(loss_partial)(coils_initial.x) - - -curvature_partial=partial(loss_coil_curvature, dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,max_coil_curvature=max_coil_curvature) -length_partial=partial(loss_coil_length, dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,max_coil_length=max_coil_length) -Baxis_average_partial=partial(loss_normB_axis_average,dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,npoints=15,target_B_on_axis=target_B_on_axis) - -# Create the constraints -penalty = 1. #Intial penalty values -multiplier=0.5 #Initial lagrange multiplier values -sq_grad=0.0 #Initial square gradient parameter value for Mu adaptative -model_lagrangian='Standard' #Use standard augmented lagragian suitable for bounded optimizers -#Since we are using LBFGS-B from jaxopt, model_mu will be updated with tolerances so we do not need to difinte the model - -#Construct constraints -constraints = alm.combine( -alm.eq(curvature_partial, multiplier=multiplier,penalty=penalty,sq_grad=sq_grad), -alm.eq(length_partial, multiplier=multiplier,penalty=penalty,sq_grad=sq_grad), -alm.eq(Baxis_average_partial, multiplier=multiplier,penalty=penalty,sq_grad=sq_grad), -alm.eq(loss_partial, multiplier=multiplier,penalty=penalty,sq_grad=sq_grad), -) - - - -beta=2. #penalty update parameter -mu_max=1.e4 #Maximum penalty parameter allowed -alpha=0.99 #These are parameters only used if gradient descent and adaaptative mu -gamma=1.e-2 -epsilon=1.e-8 -omega_tol=0.0001 #desired grad_tolerance, associated with grad of lagrangian to main parameters -eta_tol=0.001 #desired contraint tolerance, associated with variation of contraints - - - -#If loss=cost_function(x) is not prescribed, f(x)=0 is considered -ALM=alm.ALM_model_jaxopt_lbfgsb(constraints,model_lagrangian=model_lagrangian,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) - -#Initializing lagrange multipliers -lagrange_params=constraints.init(coils_initial.x) -#parameters are a tuple of the primal/main optimisation parameters and the lagrange multipliers -params = coils_initial.x, lagrange_params -#This is just to initialize an empty state for the lagrange multiplier update and get some information -lag_state,grad,info=ALM.init(params) - -#Initializing first tolerances for the inner minimisation loop iteration -mu_average=alm.penalty_average(lagrange_params) -#omega=1.#1./mu_average -#eta=1000.#1./mu_average**0.1 -omega=1./mu_average -eta=1./mu_average**0.1 - -i=0 -while i<=maximum_function_evaluations and (jnp.linalg.norm(grad[0])>omega_tol or alm.norm_constraints(info[2])>eta_tol): - #One step of ALM optimization - params, lag_state,grad,info,eta,omega = ALM.update(params,lag_state,grad,info,eta,omega) - print(f'i: {i}, loss f: {info[0]:g},loss L: {info[1]:g}, infeasibility: {alm.total_infeasibility(info[2]):g}') - #print('lagrange',params[1]) - i=i+1 - - -dofs_curves = jnp.reshape(params[0][:len_dofs_curves], (dofs_curves_shape)) -dofs_currents = params[0][len_dofs_curves:] -curves = Curves(dofs_curves, n_segments, nfp, stellsym) -new_coils = Coils(curves=curves, currents=dofs_currents*coils_initial.currents_scale) -params=new_coils.x -tracing_initial = Tracing(field=coils_initial, particles=particles, maxtime=t, model=model,times_to_trace=num_steps,timestep=timestep,boundary=boundary) -tracing_optimized = Tracing(field=new_coils, particles=particles, maxtime=t, model=model,times_to_trace=num_steps,timestep=timestep,boundary=boundary) - -#print('Final params',params) -#print(info[1]) -# Plot trajectories, before and after optimization -fig = plt.figure(figsize=(9, 8)) -ax1 = fig.add_subplot(221, projection='3d') -ax2 = fig.add_subplot(222, projection='3d') -ax3 = fig.add_subplot(223) -ax4 = fig.add_subplot(224) - -coils_initial.plot(ax=ax1, show=False) -tracing_initial.plot(ax=ax1, show=False) -for i, trajectory in enumerate(tracing_initial.trajectories): - ax3.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') - -ax3.set_xlabel('R (m)') -ax3.set_ylabel('Z (m)') -#ax3.legend() -new_coils.plot(ax=ax2, show=False) -tracing_optimized.plot(ax=ax2, show=False) -for i, trajectory in enumerate(tracing_optimized.trajectories): - ax4.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') -ax4.set_xlabel('R (m)') -ax4.set_ylabel('Z (m)')#ax4.legend() -plt.tight_layout() -plt.show() - -# # Save the coils to a json file -# coils_optimized.to_json("stellarator_coils.json") -# # Load the coils from a json file -# from essos.coils import Coils_from_json -# coils = Coils_from_json("stellarator_coils.json") - -# # Save results in vtk format to analyze in Paraview -# tracing_initial.to_vtk('trajectories_initial') -#tracing_optimized.to_vtk('trajectories_final') -#coils_initial.to_vtk('coils_initial') -#new_coils.to_vtk('coils_optimized') diff --git a/examples/optimize_coils_vmec_surface.py b/examples/optimize_coils_vmec_surface.py deleted file mode 100644 index 57324b25..00000000 --- a/examples/optimize_coils_vmec_surface.py +++ /dev/null @@ -1,81 +0,0 @@ -import os -number_of_processors_to_use = 5 # Parallelization, this should divide ntheta*nphi -os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' -from time import time -import jax.numpy as jnp -import matplotlib.pyplot as plt -from essos.surfaces import BdotN_over_B -from essos.coils import Coils, CreateEquallySpacedCurves -from essos.fields import Vmec, BiotSavart -from essos.objective_functions import loss_BdotN -from essos.optimization import optimize_loss_function - -# Optimization parameters -max_coil_length = 10 -max_coil_curvature = 1.0 -order_Fourier_series_coils = 3 -number_coil_points = order_Fourier_series_coils*15 -maximum_function_evaluations = 50 -number_coils_per_half_field_period = 3 -tolerance_optimization = 1e-5 -ntheta=35 -nphi=35 - -# Initialize VMEC field -vmec = Vmec(os.path.join(os.path.dirname(__file__), 'input_files', - 'wout_LandremanPaul2021_QA_reactorScale_lowres.nc'), - ntheta=ntheta, nphi=nphi, range_torus='half period') - -# Initialize coils -current_on_each_coil = 1 -number_of_field_periods = vmec.nfp -major_radius_coils = vmec.r_axis -minor_radius_coils = vmec.r_axis/1.8 -curves = CreateEquallySpacedCurves(n_curves=number_coils_per_half_field_period, - order=order_Fourier_series_coils, - R=major_radius_coils, r=minor_radius_coils, - n_segments=number_coil_points, - nfp=number_of_field_periods, stellsym=True) -coils_initial = Coils(curves=curves, currents=[current_on_each_coil]*number_coils_per_half_field_period) - -# Optimize coils -print(f'Optimizing coils with {maximum_function_evaluations} function evaluations.') -time0 = time() -coils_optimized = optimize_loss_function(loss_BdotN, initial_dofs=coils_initial.x, coils=coils_initial, tolerance_optimization=tolerance_optimization, - maximum_function_evaluations=maximum_function_evaluations, vmec=vmec, - max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature,) -print(f"Optimization took {time()-time0:.2f} seconds") - - -BdotN_over_B_initial = BdotN_over_B(vmec.surface, BiotSavart(coils_initial)) -BdotN_over_B_optimized = BdotN_over_B(vmec.surface, BiotSavart(coils_optimized)) -curvature=jnp.mean(BiotSavart(coils_optimized).coils.curvature, axis=1) -length=jnp.max(jnp.ravel(BiotSavart(coils_optimized).coils.length)) -print(f"Mean curvature: ",curvature) -print(f"Length:", length) -print(f"Maximum BdotN/B before optimization: {jnp.max(BdotN_over_B_initial):.2e}") -print(f"Maximum BdotN/B after optimization: {jnp.max(BdotN_over_B_optimized):.2e}") - -# Plot coils, before and after optimization -fig = plt.figure(figsize=(8, 4)) -ax1 = fig.add_subplot(121, projection='3d') -ax2 = fig.add_subplot(122, projection='3d') -coils_initial.plot(ax=ax1, show=False) -vmec.surface.plot(ax=ax1, show=False) -coils_optimized.plot(ax=ax2, show=False) -vmec.surface.plot(ax=ax2, show=False) -plt.tight_layout() -plt.show() - -# # Save the coils to a json file -# coils_optimized.to_json("stellarator_coils.json") -# # Load the coils from a json file -# from essos.coils import Coils_from_json -# coils = Coils_from_json("stellarator_coils.json") - -# # Save results in vtk format to analyze in Paraview -# from essos.fields import BiotSavart -# vmec.surface.to_vtk('surface_initial', field=BiotSavart(coils_initial)) -# vmec.surface.to_vtk('surface_final', field=BiotSavart(coils_optimized)) -# coils_initial.to_vtk('coils_initial') -# coils_optimized.to_vtk('coils_optimized') \ No newline at end of file diff --git a/examples/optimize_coils_vmec_surface_augmented_lagrangian.py b/examples/optimize_coils_vmec_surface_augmented_lagrangian.py deleted file mode 100644 index 23f08b30..00000000 --- a/examples/optimize_coils_vmec_surface_augmented_lagrangian.py +++ /dev/null @@ -1,189 +0,0 @@ -import os -number_of_processors_to_use = 1 # Parallelization, this should divide ntheta*nphi -os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' -from time import time -import jax.numpy as jnp -import matplotlib.pyplot as plt -from essos.surfaces import BdotN_over_B -from essos.coils import Coils, CreateEquallySpacedCurves,Curves -from essos.fields import Vmec, BiotSavart -from essos.objective_functions import loss_BdotN_only_constraint,loss_coil_curvature_new,loss_coil_length_new,loss_BdotN_only -from essos.objective_functions import loss_coil_curvature,loss_coil_length -from essos.objective_functions import loss_BdotN -from essos.optimization import optimize_loss_function - -import essos.augmented_lagrangian as alm -from functools import partial - -# Optimization parameters -maximum_function_evaluations=10 -max_coil_length = 40 -max_coil_curvature = 0.5 -bdotn_tol=1.e-6 -order_Fourier_series_coils = 6 -number_coil_points = order_Fourier_series_coils*10 -number_coils_per_half_field_period = 4 -ntheta=32 -nphi=32 -#Tolerance for no normal (no ALM) optimization -tolerance_optimization = 1e-5 - -# Initialize VMEC field -vmec = Vmec(os.path.join(os.path.dirname(__name__), 'input_files', - 'wout_LandremanPaul2021_QA_reactorScale_lowres.nc'), - ntheta=ntheta, nphi=nphi, range_torus='half period') - -# Initialize coils -current_on_each_coil = 1 -number_of_field_periods = vmec.nfp -major_radius_coils = vmec.r_axis -minor_radius_coils = vmec.r_axis/1.5 -curves = CreateEquallySpacedCurves(n_curves=number_coils_per_half_field_period, - order=order_Fourier_series_coils, - R=major_radius_coils, r=minor_radius_coils, - n_segments=number_coil_points, - nfp=number_of_field_periods, stellsym=True) -coils_initial = Coils(curves=curves, currents=[current_on_each_coil]*number_coils_per_half_field_period) - -len_dofs_curves = len(jnp.ravel(coils_initial.dofs_curves)) -nfp = coils_initial.nfp -stellsym = coils_initial.stellsym -n_segments = coils_initial.n_segments -dofs_curves = coils_initial.dofs_curves -currents_scale = coils_initial.currents_scale -dofs_curves_shape = coils_initial.dofs_curves.shape - - - - -# Create the constraints -penalty = 0.1 #Intial penalty values -multiplier=0.5 #Initial lagrange multiplier values -sq_grad=0.0 #Initial square gradient parameter value for Mu adaptative -model_lagrangian='Standard' #Use standard augmented lagragian suitable for bounded optimizers -#Since we are using LBFGS-B from jaxopt, model_mu will be updated with tolerances so we do not need to difinte the model - - -curvature_partial=partial(loss_coil_curvature, dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,max_coil_curvature=max_coil_curvature) -length_partial=partial(loss_coil_length, dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,max_coil_length=max_coil_length) -bdotn_partial=partial(loss_BdotN_only_constraint, vmec=vmec, dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp,n_segments=n_segments, stellsym=stellsym,target_tol=bdotn_tol) -bdotn_only_partial=partial(loss_BdotN_only, vmec=vmec, dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp,n_segments=n_segments, stellsym=stellsym) - -#Construct constraints -constraints = alm.combine( -alm.eq(curvature_partial,model_lagrangian=model_lagrangian, multiplier=multiplier,penalty=penalty,sq_grad=sq_grad), -alm.eq(length_partial,model_lagrangian=model_lagrangian, multiplier=multiplier,penalty=penalty,sq_grad=sq_grad), -alm.eq(bdotn_partial,model_lagrangian=model_lagrangian, multiplier=multiplier,penalty=penalty,sq_grad=sq_grad) -) - - - -beta=2. #penalty update parameter -mu_max=1.e4 #Maximum penalty parameter allowed -alpha=0.99 #These are parameters only used if gradient descent and adaaptative mu -gamma=1.e-2 -epsilon=1.e-8 -omega_tol=1.e-7 #desired grad_tolerance, associated with grad of lagrangian to main parameters -eta_tol=1.e-7 #desired contraint tolerance, associated with variation of contraints - - - -#If loss=cost_function(x) is not prescribed, f(x)=0 is considered -ALM=alm.ALM_model_jaxopt_lbfgsb(constraints,model_lagrangian=model_lagrangian,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) - -#Initializing lagrange multipliers -lagrange_params=constraints.init(coils_initial.x) -#parameters are a tuple of the primal/main optimisation parameters and the lagrange multipliers -params = coils_initial.x, lagrange_params -#This is just to initialize an empty state for the lagrange multiplier update and get some information -lag_state,grad,info=ALM.init(params) - -#Initializing first tolerances for the inner minimisation loop iteration -mu_average=alm.penalty_average(lagrange_params) -omega=1./mu_average -eta=1./mu_average**0.1 - - - - - -# Optimize coils -print(f'Optimizing coils with {maximum_function_evaluations} function evaluations no ALM.') -time0 = time() -coils_optimized = optimize_loss_function(loss_BdotN, initial_dofs=coils_initial.x, coils=coils_initial, tolerance_optimization=tolerance_optimization, - maximum_function_evaluations=maximum_function_evaluations, vmec=vmec, - max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature,) -print(f"Optimization took {time()-time0:.2f} seconds") - - - - - -# Optimize coils -print(f'Optimizing coils with {maximum_function_evaluations} function evaluations using ALM.') -time0 = time() - - -i=0 -while i<=maximum_function_evaluations and (jnp.linalg.norm(grad[0])>omega_tol or alm.norm_constraints(info[2])>eta_tol): - #One step of ALM optimization - params, lag_state,grad,info,eta,omega = ALM.update(params,lag_state,grad,info,eta,omega) - #if i % 5 == 0: - #print(f'i: {i}, loss f: {info[0]:g}, infeasibility: {alm.total_infeasibility(info[1]):g}') - print(f'i: {i}, loss f: {info[0]:g},loss L: {info[1]:g}, infeasibility: {alm.total_infeasibility(info[2]):g}') - #print('lagrange',params[1]) - i=i+1 - - - -dofs_curves = jnp.reshape(params[0][:len_dofs_curves], (dofs_curves_shape)) -dofs_currents = params[0][len_dofs_curves:] -curves = Curves(dofs_curves, n_segments, nfp, stellsym) -coils_optimized_alm = Coils(curves=curves, currents=dofs_currents*coils_initial.currents_scale) - -print(f"Optimization took {time()-time0:.2f} seconds") - - -BdotN_over_B_initial = BdotN_over_B(vmec.surface, BiotSavart(coils_initial)) -BdotN_over_B_optimized = BdotN_over_B(vmec.surface, BiotSavart(coils_optimized)) -curvature=jnp.mean(BiotSavart(coils_optimized).coils.curvature, axis=1) -length=jnp.max(jnp.ravel(BiotSavart(coils_optimized).coils.length)) -BdotN_over_B_optimized_alm = BdotN_over_B(vmec.surface, BiotSavart(coils_optimized_alm)) -curvature_alm=jnp.mean(BiotSavart(coils_optimized_alm).coils.curvature, axis=1) -length_alm=jnp.max(jnp.ravel(BiotSavart(coils_optimized_alm).coils.length)) - - -print(f"Maximum allowed curvature target: ",max_coil_curvature) -print(f"Maximum allowed length target: ",max_coil_length) -print(f"Mean curvature without ALM: ",curvature) -print(f"Length withou ALM:", length) -print(f"Mean curvature with ALM: ",curvature_alm) -print(f"Length with ALM:", length_alm) -print(f"Maximum BdotN/B before optimization: {jnp.max(BdotN_over_B_initial):.2e}") -print(f"Maximum BdotN/B after optimization without ALM: {jnp.max(BdotN_over_B_optimized):.2e}") -print(f"Maximum BdotN/B after optimization with ALM: {jnp.max(BdotN_over_B_optimized_alm):.2e}") -# Plot coils, before and after optimization -fig = plt.figure(figsize=(8, 4)) -ax1 = fig.add_subplot(121, projection='3d') -ax2 = fig.add_subplot(122, projection='3d') -coils_initial.plot(ax=ax1, show=False) -vmec.surface.plot(ax=ax1, show=False) -coils_optimized.plot(ax=ax2, show=False, label='Optimized no ALM') -coils_optimized_alm.plot(ax=ax2, show=False,color='orange', label='Optimized with ALM') -vmec.surface.plot(ax=ax2, show=False) -plt.legend() -plt.tight_layout() -plt.show() - -# # Save the coils to a json file -# coils_optimized.to_json("stellarator_coils.json") -# # Load the coils from a json file -# from essos.coils import Coils_from_json -# coils = Coils_from_json("stellarator_coils.json") - -# # Save results in vtk format to analyze in Paraview -# from essos.fields import BiotSavart -# vmec.surface.to_vtk('surface_initial', field=BiotSavart(coils_initial)) -# vmec.surface.to_vtk('surface_final', field=BiotSavart(coils_optimized)) -# coils_initial.to_vtk('coils_initial') -# coils_optimized.to_vtk('coils_optimized') \ No newline at end of file diff --git a/examples/optimize_coils_vmec_surface_augmented_lagrangian_stochastic.py b/examples/optimize_coils_vmec_surface_augmented_lagrangian_stochastic.py deleted file mode 100644 index fdf423a5..00000000 --- a/examples/optimize_coils_vmec_surface_augmented_lagrangian_stochastic.py +++ /dev/null @@ -1,178 +0,0 @@ -import os -number_of_processors_to_use = 8 # Parallelization, this should divide ntheta*nphi -os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' -from time import time -import jax.numpy as jnp -import matplotlib.pyplot as plt -from essos.surfaces import BdotN_over_B -from essos.coils import Coils, CreateEquallySpacedCurves,Curves -from essos.fields import Vmec, BiotSavart -from essos.objective_functions import loss_BdotN_only_constraint_stochastic,loss_coil_curvature_new,loss_coil_length_new,loss_BdotN_only_stochastic -from essos.objective_functions import loss_coil_curvature,loss_coil_length -from essos.coil_perturbation import GaussianSampler - -import essos.augmented_lagrangian as alm -from functools import partial - -# Optimization parameters -maximum_function_evaluations=10 -max_coil_length = 40 -max_coil_curvature = 0.5 -bdotn_tol=1.e-6 -order_Fourier_series_coils = 6 -number_coil_points = order_Fourier_series_coils*10 -number_coils_per_half_field_period = 4 -ntheta=32 -nphi=32 - - - - - -# Initialize VMEC field -vmec = Vmec(os.path.join(os.path.dirname(__name__), 'input_files', - 'wout_LandremanPaul2021_QA_reactorScale_lowres.nc'), - ntheta=ntheta, nphi=nphi, range_torus='full torus') - -# Initialize coils -current_on_each_coil = 1 -number_of_field_periods = vmec.nfp -major_radius_coils = vmec.r_axis -minor_radius_coils = vmec.r_axis/1.5 -curves = CreateEquallySpacedCurves(n_curves=number_coils_per_half_field_period, - order=order_Fourier_series_coils, - R=major_radius_coils, r=minor_radius_coils, - n_segments=number_coil_points, - nfp=number_of_field_periods, stellsym=True) -coils_initial = Coils(curves=curves, currents=[current_on_each_coil]*number_coils_per_half_field_period) - -len_dofs_curves = len(jnp.ravel(coils_initial.dofs_curves)) -nfp = coils_initial.nfp -stellsym = coils_initial.stellsym -n_segments = coils_initial.n_segments -dofs_curves = coils_initial.dofs_curves -currents_scale = coils_initial.currents_scale -dofs_curves_shape = coils_initial.dofs_curves.shape - - - -#Sampling parameters -sigma=0.01 -length_scale=0.4*jnp.pi -n_derivs=2 -N_samples=10 #Number of samples for the stochastic perturbation -#Create a Gaussian sampler for perturbation -#This sampler will be used to perturb the coils -sampler=GaussianSampler(coils_initial.quadpoints,sigma=sigma,length_scale=length_scale,n_derivs=n_derivs) - - - - -# Create the constraints -penalty = 0.1 #Intial penalty values -multiplier=0.5 #Initial lagrange multiplier values -sq_grad=0.0 #Initial square gradient parameter value for Mu adaptative -model_lagrangian='Standard' #Use standard augmented lagragian suitable for bounded optimizers -#Since we are using LBFGS-B from jaxopt, model_mu will be updated with tolerances so we do not need to difinte the model - - -curvature_partial=partial(loss_coil_curvature, dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,max_coil_curvature=max_coil_curvature) -length_partial=partial(loss_coil_length, dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym,max_coil_length=max_coil_length) -bdotn_partial=partial(loss_BdotN_only_constraint_stochastic,sampler=sampler,N_samples=N_samples, vmec=vmec, dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp,n_segments=n_segments, stellsym=stellsym,target_tol=bdotn_tol) -bdotn_only_partial=partial(loss_BdotN_only_stochastic,sampler=sampler,N_samples=N_samples, vmec=vmec, dofs_curves=coils_initial.dofs_curves, currents_scale=currents_scale, nfp=nfp,n_segments=n_segments, stellsym=stellsym) - -#Construct constraints -constraints = alm.combine( -alm.eq(curvature_partial,model_lagrangian=model_lagrangian, multiplier=multiplier,penalty=penalty,sq_grad=sq_grad), -alm.eq(length_partial,model_lagrangian=model_lagrangian, multiplier=multiplier,penalty=penalty,sq_grad=sq_grad), -alm.eq(bdotn_partial,model_lagrangian=model_lagrangian, multiplier=multiplier,penalty=penalty,sq_grad=sq_grad) -) - - - -beta=2. #penalty update parameter -mu_max=1.e4 #Maximum penalty parameter allowed -alpha=0.99 #These are parameters only used if gradient descent and adaaptative mu -gamma=1.e-2 -epsilon=1.e-8 -omega_tol=1.e-7 #desired grad_tolerance, associated with grad of lagrangian to main parameters -eta_tol=1.e-7 #desired contraint tolerance, associated with variation of contraints - - - -#If loss=cost_function(x) is not prescribed, f(x)=0 is considered -ALM=alm.ALM_model_jaxopt_lbfgsb(constraints,model_lagrangian=model_lagrangian,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) - -#Initializing lagrange multipliers -lagrange_params=constraints.init(coils_initial.x) -#parameters are a tuple of the primal/main optimisation parameters and the lagrange multipliers -params = coils_initial.x, lagrange_params -#This is just to initialize an empty state for the lagrange multiplier update and get some information -lag_state,grad,info=ALM.init(params) - -#Initializing first tolerances for the inner minimisation loop iteration -mu_average=alm.penalty_average(lagrange_params) -omega=1./mu_average -eta=1./mu_average**0.1 - - - - -# Optimize coils -print(f'Optimizing coils with {maximum_function_evaluations} function evaluations using stochastic and ALM.') -time0 = time() - - -i=0 -while i<=maximum_function_evaluations and (jnp.linalg.norm(grad[0])>omega_tol or alm.norm_constraints(info[2])>eta_tol): - #One step of ALM optimization - params, lag_state,grad,info,eta,omega = ALM.update(params,lag_state,grad,info,eta,omega) - #if i % 5 == 0: - #print(f'i: {i}, loss f: {info[0]:g}, infeasibility: {alm.total_infeasibility(info[1]):g}') - print(f'i: {i}, loss f: {info[0]:g},loss L: {info[1]:g}, infeasibility: {alm.total_infeasibility(info[2]):g}') - #print('lagrange',params[1]) - i=i+1 - - - -dofs_curves = jnp.reshape(params[0][:len_dofs_curves], (dofs_curves_shape)) -dofs_currents = params[0][len_dofs_curves:] -curves = Curves(dofs_curves, n_segments, nfp, stellsym) -coils_optimized = Coils(curves=curves, currents=dofs_currents*coils_initial.currents_scale) - -print(f"Stochastic optimization with ALM took {time()-time0:.2f} seconds") - - -BdotN_over_B_initial = BdotN_over_B(vmec.surface, BiotSavart(coils_initial)) -BdotN_over_B_optimized = BdotN_over_B(vmec.surface, BiotSavart(coils_optimized)) -curvature=jnp.mean(BiotSavart(coils_optimized).coils.curvature, axis=1) -length=jnp.max(jnp.ravel(BiotSavart(coils_optimized).coils.length)) -print(f"Mean curvature: ",curvature) -print(f"Length:", length) -print(f"Maximum BdotN/B before optimization: {jnp.max(BdotN_over_B_initial):.2e}") -print(f"Maximum BdotN/B after optimization: {jnp.max(BdotN_over_B_optimized):.2e}") -print(f"Average BdotN/B before optimization: {jnp.average(jnp.absolute(BdotN_over_B_initial)):.2e}") -print(f"Average BdotN/B after optimization: {jnp.average(jnp.absolute(BdotN_over_B_optimized)):.2e}") -# Plot coils, before and after optimization -fig = plt.figure(figsize=(8, 4)) -ax1 = fig.add_subplot(121, projection='3d') -ax2 = fig.add_subplot(122, projection='3d') -coils_initial.plot(ax=ax1, show=False) -vmec.surface.plot(ax=ax1, show=False) -coils_optimized.plot(ax=ax2, show=False) -vmec.surface.plot(ax=ax2, show=False) -plt.tight_layout() -plt.show() - -# # Save the coils to a json file -# coils_optimized.to_json("stellarator_coils.json") -# # Load the coils from a json file -# from essos.coils import Coils_from_json -# coils = Coils_from_json("stellarator_coils.json") - -# # Save results in vtk format to analyze in Paraview -# from essos.fields import BiotSavart -# vmec.surface.to_vtk('surface_initial', field=BiotSavart(coils_initial)) -# vmec.surface.to_vtk('surface_final', field=BiotSavart(coils_optimized)) -# coils_initial.to_vtk('coils_initial') -# coils_optimized.to_vtk('coils_optimized') \ No newline at end of file diff --git a/examples/paper/fo_integrators.py b/examples/paper/fo_integrators.py new file mode 100644 index 00000000..dc652501 --- /dev/null +++ b/examples/paper/fo_integrators.py @@ -0,0 +1,95 @@ +import os +number_of_processors_to_use = 1 # Parallelization, this should divide nparticles +os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' +from time import time +from jax import block_until_ready +import jax.numpy as jnp +import matplotlib.pyplot as plt +plt.rcParams.update({'font.size': 18}) +from essos.fields import BiotSavart +from essos.coils import Coils +from essos.constants import PROTON_MASS, ONE_EV, ELEMENTARY_CHARGE +from essos.dynamics import Tracing, Particles +import diffrax + +output_dir = os.path.join(os.path.dirname(__file__), 'output') +if not os.path.exists(output_dir): + os.makedirs(output_dir) + +# Load coils and field +json_file = os.path.join(os.path.dirname(__file__), '../input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') +coils = Coils.from_json(json_file) +field = BiotSavart(coils) + +# Particle parameters +nparticles = number_of_processors_to_use +mass=PROTON_MASS +energy=5000*ONE_EV +cyclotron_frequency = ELEMENTARY_CHARGE*0.3/mass +print("cyclotron period:", 1/cyclotron_frequency) + +# Particles initialization +initial_xyz=jnp.array([[1.23, 0, 0]]) +particles = Particles(initial_xyz=initial_xyz, mass=mass, energy=energy, initial_vparallel_over_v=[0.8], field=field) + +# Tracing parameters +tmax = 1e-4 + +fig, ax = plt.subplots(figsize=(9, 6)) + +# Adaptive Runge-Kutta solvers: for each solver, sweep the integration tolerance +# and record (computation time, energy error). Tighter tolerance -> smaller error +# but longer runtime, which traces out each curve in the time-vs-error plot. +adaptive_solvers = [('Tsit5', diffrax.Tsit5), ('Dopri5', diffrax.Dopri5), ('Dopri8', diffrax.Dopri8)] +for method_name, solver_class in adaptive_solvers: + energies = [] + tracing_times = [] + for trace_tolerance in [1e-8, 1e-9, 1e-10, 1e-11, 1e-12, 1e-13, 1e-14, 1e-15]: + time0 = time() + tracing = Tracing(field=field, model='FullOrbit', particles=particles, + maxtime=tmax, timestep=1e-9, + atol=trace_tolerance, rtol=trace_tolerance, + solver=solver_class()) + block_until_ready(tracing.trajectories) + tracing_times += [time() - time0] + + print(f"Tracing with adaptive {method_name} and tolerance {trace_tolerance:.0e} took {tracing_times[-1]:.2f} seconds") + + energies += [jnp.mean(jnp.abs(tracing.energy()-particles.energy)/particles.energy)] + ax.plot(tracing_times, energies, label=f'{method_name} adapt', marker='o', markersize=3, linestyle='-') + +# Boris pusher: a fixed-step symplectic integrator. Instead of a tolerance, we +# sweep the number of timesteps per gyration (more points per orbit -> smaller +# step -> smaller error), which traces out its own time-vs-error curve. +energies = [] +tracing_times = [] +for n_points_in_gyration in [10, 20, 50, 75, 100, 150, 200]: + dt = 1/(n_points_in_gyration*cyclotron_frequency) + time0 = time() + tracing = Tracing(field=field, model='FullOrbit_Boris', particles=particles, + maxtime=tmax, timestep=dt) + block_until_ready(tracing.trajectories) + tracing_times += [time() - time0] + + print(f"Tracing with Boris and step {dt:.2e} took {tracing_times[-1]:.2f} seconds") + + energies += [jnp.mean(jnp.abs(tracing.energy()-particles.energy)/particles.energy)] +ax.plot(tracing_times, energies, label='Boris', marker='o', markersize=4, linestyle='-') + + +ax.legend(fontsize=15, loc='upper left') +ax.set_xlabel('Computation time (s)') +ax.set_ylabel('Relative energy error') +ax.set_xscale('log') +ax.set_yscale('log') +ax.set_xlim(1e-1, 1e2) +ax.set_ylim(1e-16, 1e-4) +plt.grid(axis='x', which='both', linestyle='--', linewidth=0.6) +plt.grid(axis='y', which='major', linestyle='--', linewidth=0.6) +plt.tight_layout() +plt.savefig(os.path.join(output_dir, 'fo_integration.pdf')) +plt.show() + +## Save results in vtk format to analyze in Paraview +# tracing.to_vtk('trajectories') +# coils.to_vtk('coils') diff --git a/examples/paper/gc_integrators.py b/examples/paper/gc_integrators.py new file mode 100644 index 00000000..cac018db --- /dev/null +++ b/examples/paper/gc_integrators.py @@ -0,0 +1,95 @@ +import os +import gc +number_of_processors_to_use = 1 # Parallelization, this should divide nparticles +os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' +from time import time +from jax import block_until_ready +import jax.numpy as jnp +import matplotlib.pyplot as plt +plt.rcParams.update({'font.size': 18}) +from essos.fields import BiotSavart +from essos.coils import Coils +from essos.constants import PROTON_MASS, ONE_EV, ELEMENTARY_CHARGE +from essos.dynamics import Tracing, Particles +import diffrax + +output_dir = os.path.join(os.path.dirname(__file__), 'output') +if not os.path.exists(output_dir): + os.makedirs(output_dir) + +# Load coils and field +json_file = os.path.join(os.path.dirname(__file__), '../input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') +coils = Coils.from_json(json_file) +field = BiotSavart(coils) + +# Particle parameters +nparticles = number_of_processors_to_use +mass=PROTON_MASS +energy=5000*ONE_EV +cyclotron_frequency = ELEMENTARY_CHARGE*0.3/mass +print("cyclotron period:", 1/cyclotron_frequency) + +# Particles initialization +initial_xyz=jnp.array([[1.23, 0, 0]]) +particles = Particles(initial_xyz=initial_xyz, mass=mass, energy=energy, initial_vparallel_over_v=[0.8], field=field) + +# Tracing parameters +tmax = 1e-4 + +# Two figures: energy error vs computation time, and energy error vs tolerance. +fig, ax = plt.subplots(figsize=(9, 6)) +fig_tol, ax_tol = plt.subplots(figsize=(9, 6)) +markers = ["o-", "^-", "*-", "s-"] +# For each adaptive solver, sweep the integration tolerance and record both the +# resulting energy error and the wall-clock time. Kvaerno5 is implicit; the +# others are explicit Runge-Kutta methods. +solvers = [('Tsit5', diffrax.Tsit5), + ('Dopri5', diffrax.Dopri5), + ('Dopri8', diffrax.Dopri8), + ('Kvaerno5', diffrax.Kvaerno5)] +for (method, solver_class), marker in zip(solvers, markers): + energies = [] + tracing_times = [] + tolerances = [1e-7, 1e-8, 1e-9, 1e-10, 1e-11, 1e-12, 1e-13, 1e-14, 1e-15, 1e-16] + for tolerance in tolerances: + time0 = time() + tracing = Tracing(field=field, model='GuidingCenter', particles=particles, + maxtime=tmax, timestep=1e-7, + atol=tolerance, rtol=tolerance, + solver=solver_class()) + block_until_ready(tracing.trajectories) + tracing_times += [time() - time0] + + print(f"Tracing with adaptive {method} and {tolerance=:.0e} took {tracing_times[-1]:.2f} seconds") + + energies += [jnp.max(jnp.abs(tracing.energy()-particles.energy)/particles.energy)] + ax.plot(tracing_times, energies, label=f'{method} adapt', marker='o', markersize=3) + ax_tol.plot(tolerances, energies, marker, label=f'{method} adapt', clip_on=False, linewidth=2.5) + gc.collect() + +ax.set_xlabel('Computation time (s)') +ax_tol.set_xlabel('Tracing tolerance') +ax.set_xlim(1e-1, 1e2) +ax_tol.set_xlim(tolerances[-1], tolerances[0]) + +for axis in [ax, ax_tol]: + axis.legend(fontsize=15) + axis.set_ylabel('Relative energy error') + axis.set_xscale('log') + axis.set_yscale('log') + axis.set_ylim(1e-16, 1e-4) + axis.grid(axis='x', which='both', linestyle='--', linewidth=0.6) + axis.grid(axis='y', which='major', linestyle='--', linewidth=0.6) +for figure in [fig, fig_tol]: + figure.tight_layout() + +for spine in ax_tol.spines.values(): + spine.set_zorder(0) + +fig.savefig(os.path.join(output_dir, 'gc_integration.pdf')) +fig_tol.savefig(os.path.join(output_dir, 'energy_vs_tol.pdf')) +plt.show() + +## Save results in vtk format to analyze in Paraview +# tracing.to_vtk('trajectories') +# coils.to_vtk('coils') diff --git a/examples/paper/gc_vs_fo.py b/examples/paper/gc_vs_fo.py new file mode 100644 index 00000000..216a9b2e --- /dev/null +++ b/examples/paper/gc_vs_fo.py @@ -0,0 +1,87 @@ +import os +number_of_processors_to_use = 1 # Parallelization, this should divide nparticles +os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' +from jax import vmap +from time import time +import jax.numpy as jnp +import matplotlib.pyplot as plt +plt.rcParams.update({'font.size': 18}) +from essos.fields import BiotSavart +from essos.coils import Coils +from essos.constants import PROTON_MASS, ONE_EV, ELEMENTARY_CHARGE +from essos.dynamics import Tracing, Particles +from jax import block_until_ready + +output_dir = os.path.join(os.path.dirname(__file__), 'output') +if not os.path.exists(output_dir): + os.makedirs(output_dir) + +# Load coils and field +json_file = os.path.join(os.path.dirname(__file__), '../input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') +coils = Coils.from_json(json_file) +field = BiotSavart(coils) + +# Particle parameters +nparticles = number_of_processors_to_use +mass=PROTON_MASS +energy=5000*ONE_EV +cyclotron_frequency = ELEMENTARY_CHARGE*0.3/mass +print("cyclotron period:", 1/cyclotron_frequency) + +# Particles initialization +initial_xyz=jnp.array([[1.23, 0, 0]]) + +particles_passing = Particles(initial_xyz=initial_xyz, mass=mass, energy=energy, initial_vparallel_over_v=[0.1], phase_angle_full_orbit=0) +particles_traped = Particles(initial_xyz=initial_xyz, mass=mass, energy=energy, initial_vparallel_over_v=[0.9], phase_angle_full_orbit=0) +particles = particles_passing.join(particles_traped, field=field) + +# Tracing parameters +tmax = 1e-5 +trace_tolerance = 1e-14 +dt_gc = 1e-7 +dt_fo = 1e-9 +num_steps_gc = int(tmax/dt_gc) +num_steps_fo = int(tmax/dt_fo) + +# Trace in ESSOS +time0 = time() +tracing_gc = Tracing(field=field, model='GuidingCenter', particles=particles, + maxtime=tmax, timestep=dt_gc, atol=trace_tolerance, rtol=trace_tolerance, + times_to_trace=200) +trajectories_guidingcenter = block_until_ready(tracing_gc.trajectories) +print(f"ESSOS guiding center tracing took {time()-time0:.2f} seconds") + +time0 = time() +tracing_fo = Tracing(field=field, model='FullOrbit', particles=particles, maxtime=tmax, + timestep=dt_fo, atol=trace_tolerance, rtol=trace_tolerance, + times_to_trace=600) + +block_until_ready(tracing_fo.trajectories) +print(f"ESSOS full orbit tracing took {time()-time0:.2f} seconds") + +# Plot trajectories, velocity parallel to the magnetic field, and energy error +fig = plt.figure(figsize=(9, 8)) +ax = fig.add_subplot(projection='3d') +coils.plot(ax=ax, show=False) +tracing_gc.plot(ax=ax, show=False, color='black', linewidth=2) +tracing_fo.plot(ax=ax, show=False) +plt.tight_layout() + +plt.figure(figsize=(9, 6)) +plt.plot(tracing_gc.times[1:]*1000, jnp.abs(tracing_gc.energy()[0][1:]/particles.energy-1)+1e-17, label='Guiding Center', color='red') +plt.plot(tracing_fo.times[1:]*1000, jnp.abs(tracing_fo.energy()[0][1:]/particles.energy-1)+1e-17, label='Full Orbit', color='blue') +plt.xlabel('Time (ms)') +plt.ylabel('Relative energy error') +plt.xlim(0, tmax*1000) +# plt.ylim(bottom=0) +plt.yscale('log') +plt.legend() +plt.tight_layout() +plt.savefig(os.path.join(output_dir, 'energies.png'), dpi=300) + +plt.show() + +## Save results in vtk format to analyze in Paraview +# tracing_gc.to_vtk(os.path.join(output_dir, 'trajectories_gc')) +# tracing_fo.to_vtk(os.path.join(output_dir, 'trajectories_fo')) +# coils.to_vtk(os.path.join(output_dir, 'coils')) \ No newline at end of file diff --git a/examples/paper/gradients.py b/examples/paper/gradients.py new file mode 100644 index 00000000..1b946349 --- /dev/null +++ b/examples/paper/gradients.py @@ -0,0 +1,110 @@ +import os +import jax.numpy as jnp +from jax import block_until_ready +from time import perf_counter as timer +import matplotlib.pyplot as plt +plt.rcParams.update({'font.size': 18}) + +from essos.coils import Coils, CreateEquallySpacedCurves +from essos.fields import Vmec +from essos.fields import BiotSavart +from essos.surfaces import BdotN +from essos.losses import custom_loss + +output_dir = os.path.join(os.path.dirname(__file__), 'output') +if not os.path.exists(output_dir): + os.makedirs(output_dir) + +# Initialize VMEC field +input_filepath = os.path.join(os.path.dirname(__file__), "..", "input_files") +vmec_input = os.path.join(input_filepath, 'wout_LandremanPaul2021_QA_reactorScale_lowres.nc') +vmec = Vmec(vmec_input, ntheta=32, nphi=32, range_torus='half period') + +# Initialize coils +FOURIER_ORDER = 6 +N_SEGMENTS = FOURIER_ORDER*10 +N_COILS = 4 +COIL_CURRENT = 1. +NFP = vmec.nfp +STELLSYM = True +LARGE_R = vmec.r_axis +SMALL_R = vmec.r_axis/1.5 + +curves = CreateEquallySpacedCurves(N_COILS, FOURIER_ORDER, LARGE_R, SMALL_R, n_segments=N_SEGMENTS, nfp=NFP, stellsym=STELLSYM) +coils = Coils(curves=curves, currents=[COIL_CURRENT]*N_COILS) +field = BiotSavart(coils) + +""" Creating the loss functions """ +def loss(field, surface): + return jnp.sum(jnp.abs(BdotN(surface, field))) + +Loss = custom_loss(loss, "field", surface=vmec.surface) +Loss.dependencies = {"field": field} +dofs = Loss.starting_dofs + +loss_value = Loss(dofs) +grad_loss = Loss.grad(dofs) +print("Loss value:", loss_value) +print("Gradient:", grad_loss) + +t_start = timer() +block_until_ready(Loss(dofs)) +t_end = timer() +print(f"Loss took {t_end - t_start:.4f} seconds. Gradient would take {(t_end - t_start)*(coils.x.size +1):.4f} seconds") + +t_start = timer() +block_until_ready(Loss.grad(dofs)) +t_end = timer() +print(f"Gradient took {t_end - t_start:.4f} seconds") + +# Parameter to perturb +param = 42 + +# Set the possible perturbations +h_list = jnp.arange(-9, -0.9, 1/3) +h_list = 10.**h_list + +# Number of orders for finite differences +fd_loss = jnp.zeros(4) + +# Array to store the relative difference +fd_diff = jnp.zeros((fd_loss.size, h_list.size)) + +# Compute finite differences +for index, h in enumerate(h_list): + delta = jnp.zeros_like(dofs) + delta = delta.at[param].set(h) + + # 1st order finite differences + fd_loss = fd_loss.at[0].set((Loss(dofs+delta)-Loss(dofs))/h) + # 2nd order finite differences + fd_loss = fd_loss.at[1].set((Loss(dofs+delta)-Loss(dofs-delta))/(2*h)) + # 4th order finite differences + fd_loss = fd_loss.at[2].set((Loss(dofs-2*delta)-8*Loss(dofs-delta)+8*Loss(dofs+delta)-Loss(dofs+2*delta))/(12*h)) + # 6th order finite differences + fd_loss = fd_loss.at[3].set((Loss(dofs+3*delta)-9*Loss(dofs+2*delta)+45*Loss(dofs+delta)-45*Loss(dofs-delta)+9*Loss(dofs-2*delta)-Loss(dofs-3*delta))/(60*h)) + + fd_diff_h = jnp.abs((grad_loss[param]-fd_loss)/grad_loss[param]) + fd_diff = fd_diff.at[:, index].set(fd_diff_h) + + +# plot relative difference +plt.figure(figsize=(9, 6)) +plt.plot(h_list, fd_diff[0], "o-", label=f'1st order', clip_on=False, linewidth=2.5) +plt.plot(h_list, fd_diff[1], "^-", label=f'2nd order', clip_on=False, linewidth=2.5) +plt.plot(h_list, fd_diff[2], "*-", label=f'4th order', clip_on=False, linewidth=2.5) +plt.plot(h_list, fd_diff[3], "s-", label=f'6th order', clip_on=False, linewidth=2.5) +plt.legend(fontsize=15) +plt.xlabel('Finite differences stepsize h') +plt.ylabel('Relative error') +plt.xscale('log') +plt.yscale('log') +plt.ylim(1e-13, 1e-1) +plt.xlim(jnp.min(h_list), jnp.max(h_list)) +plt.grid(which='both', axis='x', linestyle='--', linewidth=0.6) +plt.grid(which='major', axis='y', linestyle='--', linewidth=0.6) +for spine in plt.gca().spines.values(): + spine.set_zorder(0) +plt.tight_layout() +plt.savefig(os.path.join(output_dir, 'gradients.pdf')) +plt.show() \ No newline at end of file diff --git a/examples/paper/poincare_plots.py b/examples/paper/poincare_plots.py new file mode 100644 index 00000000..cdc411a2 --- /dev/null +++ b/examples/paper/poincare_plots.py @@ -0,0 +1,134 @@ +import os +from functools import partial +number_of_processors_to_use = 1 # Parallelization, this should divide ntheta*nphi +os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' +from time import time +from jax import jit, grad, block_until_ready +import jax.numpy as jnp +import matplotlib.pyplot as plt +plt.rcParams.update({'font.size': 18}) +from essos.coils import Coils +from essos.constants import PROTON_MASS, ONE_EV, ELEMENTARY_CHARGE +from essos.fields import BiotSavart +from essos.dynamics import Tracing, Particles + +output_dir = os.path.join(os.path.dirname(__file__), 'output') +if not os.path.exists(output_dir): + os.makedirs(output_dir) + +# Input parameters +tmax_fl = 50000 +tmax_gc = 1e-3 +tmax_fo = 1e-3 + +nparticles = number_of_processors_to_use*1 +nfieldlines = number_of_processors_to_use*8 +s = 0.25 # s-coordinate: flux surface label +trace_tolerance = 1e-15 +dt_fo = 1e-9 +dt_gc = 1e-7 +timesteps_gc = int(tmax_gc/dt_gc) +timesteps_fo = int(tmax_fo/dt_fo) +mass = PROTON_MASS +energy = 4000*ONE_EV +print("cyclotron period:", 1/(ELEMENTARY_CHARGE*0.3/mass)) + +# Load coils and field +json_file = os.path.join(os.path.dirname(__file__), '../input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') +coils = Coils.from_json(json_file) +field = BiotSavart(coils) + +R0_fieldlines = jnp.linspace(1.21, 1.41, nfieldlines) +R0_particles= jnp.linspace(1.21, 1.41, nparticles) +Z0_fieldlines = jnp.zeros(nfieldlines) +Z0_particles = jnp.zeros(nparticles) +phi0_fieldlines = jnp.zeros(nfieldlines) +phi0_particles = jnp.zeros(nparticles) + +initial_xyz_fieldlines=jnp.array([R0_fieldlines*jnp.cos(phi0_fieldlines), R0_fieldlines*jnp.sin(phi0_fieldlines), Z0_fieldlines]).T +initial_xyz_particles=jnp.array([R0_particles*jnp.cos(phi0_particles), R0_particles*jnp.sin(phi0_particles), Z0_particles]).T + +particles = Particles(initial_xyz=initial_xyz_particles, mass=mass, energy=energy, field=field, min_vparallel_over_v=0.8) + +# Trace in ESSOS +# time0 = time() +# tracing_fl = Tracing(field=field, model='FieldLine', initial_conditions=initial_xyz_fieldlines, +# maxtime=tmax_fl, timesteps=tmax_fl*10, tol_step_size=trace_tolerance) +# block_until_ready(tracing_fl) +# print(f"ESSOS tracing of {nfieldlines} field lines took {time()-time0:.2f} seconds") + +time0 = time() +tracing_fo = Tracing(field=field, model='FullOrbit', particles=particles, maxtime=tmax_fo, + timestep=dt_fo, atol=trace_tolerance, rtol=trace_tolerance) +# tracing_fo.trajectories = tracing_fo.trajectories[:, 0::100, :] +# tracing_fo.times = tracing_fo.times[0::100] +# tracing_fo.energy = tracing_fo.energy[:, 0::100] +block_until_ready(tracing_fo) +print(f"ESSOS tracing of {nparticles} particles with FO for {tmax_fo:.1e}s took {time()-time0:.2f} seconds") + +time0 = time() +tracing_gc = Tracing(field=field, model='GuidingCenter', particles=particles, maxtime=tmax_gc, + timestep=dt_gc, atol=trace_tolerance, rtol=trace_tolerance) +block_until_ready(tracing_gc) +print(f"ESSOS tracing of {nparticles} particles with GC for {tmax_gc:.1e}s took {time()-time0:.2f} seconds") + +# fig = plt.figure(figsize=(9, 6)) +# ax = fig.add_subplot(projection='3d') +# coils.plot(ax=ax, show=False) +# tracing_fl.plot(ax=ax, show=False) +# plt.tight_layout() + +# fig = plt.figure(figsize=(9, 6)) +# ax = fig.add_subplot(projection='3d') +# coils.plot(ax=ax, show=False) +# tracing_fo.plot(ax=ax, show=False) +# plt.tight_layout() + +# fig = plt.figure(figsize=(9, 6)) +# ax = fig.add_subplot(projection='3d') +# coils.plot(ax=ax, show=False) +# tracing_gc.plot(ax=ax, show=False) +# plt.tight_layout() + +# fig, ax = plt.subplots(figsize=(9, 6)) +# time0 = time() +# tracing_fl.poincare_plot(ax=ax, shifts=[jnp.pi/2], show=False, s=0.5) +# print(f"ESSOS Poincare plot of {nfieldlines} field lines took {time()-time0:.2f} seconds") +# plt.xlabel('R (m)') +# plt.ylabel('Z (m)') +# ax.set_xlim(0.3, 1.3) +# ax.set_ylim(-0.3, 0.3) +# plt.grid(visible=False) +# plt.tight_layout() +# plt.savefig(os.path.join(output_dir, 'poincare_plot_fl.png'), dpi=300) +# plt.savefig(os.path.join(os.path.dirname(__file__), 'poincare_plot_fl.png'), dpi=300) + + +# fig, ax = plt.subplots(figsize=(9, 6)) +# time0 = time() +# tracing_fo.poincare_plot(ax=ax, shifts=[jnp.pi/2], show=False) +# print(f"ESSOS Poincare plot of {nparticles} particles took {time()-time0:.2f} seconds") +# plt.xlabel('R (m)') +# plt.ylabel('Z (m)') +# plt.xlim(0.3, 1.3) +# plt.ylim(-0.3, 0.3) +# plt.grid(visible=False) +# plt.tight_layout() +# plt.savefig(os.path.join(output_dir 'poincare_plot_fo.png'), dpi=300) +# plt.savefig(os.path.join(os.path.dirname(__file__), 'poincare_plot_fo.png'), dpi=300) + + +# fig, ax = plt.subplots(figsize=(9, 6)) +# time0 = time() +# tracing_gc.poincare_plot(ax=ax, shifts=[jnp.pi/2], show=False) +# print(f"ESSOS Poincare plot of {nparticles} particles took {time()-time0:.2f} seconds") +# plt.xlabel('R (m)') +# plt.ylabel('Z (m)') +# ax.set_xlim(0.3, 1.3) +# ax.set_ylim(-0.3, 0.3) +# plt.grid(visible=False) +# plt.tight_layout() +# plt.savefig(os.path.join(output_dir, 'poincare_plot_gc.png'), dpi=300) +# plt.savefig(os.path.join(os.path.dirname(__file__), 'poincare_plot_gc.png'), dpi=300) + +# plt.show() \ No newline at end of file diff --git a/examples/trace_particles_coils_fullorbit.py b/examples/particle_tracing/trace_particles_coils_fullorbit.py similarity index 88% rename from examples/trace_particles_coils_fullorbit.py rename to examples/particle_tracing/trace_particles_coils_fullorbit.py index baa59760..06abf8bc 100644 --- a/examples/trace_particles_coils_fullorbit.py +++ b/examples/particle_tracing/trace_particles_coils_fullorbit.py @@ -6,7 +6,7 @@ import jax.numpy as jnp import matplotlib.pyplot as plt from essos.fields import BiotSavart -from essos.coils import Coils_from_json +from essos.coils import Coils from essos.constants import PROTON_MASS, ONE_EV from essos.dynamics import Tracing, Particles @@ -22,8 +22,8 @@ energy=4000*ONE_EV # Load coils and field -json_file = os.path.join(os.path.dirname(__file__), 'input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') -coils = Coils_from_json(json_file) +json_file = os.path.join(os.path.dirname(__file__), '..', 'input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') +coils = Coils.from_json(json_file) field = BiotSavart(coils) # Initialize particles @@ -50,7 +50,7 @@ tracing.plot(ax=ax1, show=False) for i, trajectory in enumerate(trajectories): - ax2.plot(tracing.times, jnp.abs(tracing.energy[i]-particles.energy)/particles.energy, label=f'Particle {i+1}', linewidth=0.2) + ax2.plot(tracing.times, jnp.abs(tracing.energy()[i]-particles.energy)/particles.energy, label=f'Particle {i+1}', linewidth=0.2) def compute_v_parallel(trajectory_t): magnetic_field_unit_vector = field.B(trajectory_t[:3]) / field.AbsB(trajectory_t[:3]) return jnp.dot(trajectory_t[3:], magnetic_field_unit_vector) diff --git a/examples/trace_particles_coils_guidingcenter.py b/examples/particle_tracing/trace_particles_coils_guidingcenter.py similarity index 88% rename from examples/trace_particles_coils_guidingcenter.py rename to examples/particle_tracing/trace_particles_coils_guidingcenter.py index 018317c3..b7a52a5a 100644 --- a/examples/trace_particles_coils_guidingcenter.py +++ b/examples/particle_tracing/trace_particles_coils_guidingcenter.py @@ -5,7 +5,7 @@ import jax.numpy as jnp import matplotlib.pyplot as plt from essos.fields import BiotSavart -from essos.coils import Coils_from_json +from essos.coils import Coils from essos.constants import ALPHA_PARTICLE_MASS, ALPHA_PARTICLE_CHARGE, ONE_EV from essos.dynamics import Tracing, Particles @@ -21,8 +21,8 @@ energy=4000*ONE_EV # Load coils and field -json_file = os.path.join(os.path.dirname(__name__), 'input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') -coils = Coils_from_json(json_file) +json_file = os.path.join(os.path.dirname(__file__), '..', 'input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') +coils = Coils.from_json(json_file) field = BiotSavart(coils) # Initialize particles @@ -49,7 +49,7 @@ tracing.plot(ax=ax1, show=False) for i, trajectory in enumerate(trajectories): - ax2.plot(tracing.times, jnp.abs(tracing.energy[i]-particles.energy)/particles.energy, label=f'Particle {i+1}') + ax2.plot(tracing.times, jnp.abs(tracing.energy()[i]-particles.energy)/particles.energy, label=f'Particle {i+1}') ax3.plot(tracing.times, trajectory[:, 3]/particles.total_speed, label=f'Particle {i+1}') ax4.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') ax2.set_xlabel('Time (s)') diff --git a/examples/trace_particles_coils_guidingcenter_with_classifier.py b/examples/particle_tracing/trace_particles_coils_guidingcenter_with_classifier.py similarity index 67% rename from examples/trace_particles_coils_guidingcenter_with_classifier.py rename to examples/particle_tracing/trace_particles_coils_guidingcenter_with_classifier.py index 0a842e8d..aa05ff7c 100644 --- a/examples/trace_particles_coils_guidingcenter_with_classifier.py +++ b/examples/particle_tracing/trace_particles_coils_guidingcenter_with_classifier.py @@ -1,15 +1,18 @@ import os number_of_processors_to_use = 1 # Parallelization, this should divide nparticles os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' +import jax +print(jax.devices()) from time import time from jax import block_until_ready import jax.numpy as jnp import matplotlib.pyplot as plt from essos.fields import BiotSavart,Vmec from essos.surfaces import SurfaceClassifier -from essos.coils import Coils_from_simsopt +from essos.coils import Coils from essos.constants import ALPHA_PARTICLE_MASS, ALPHA_PARTICLE_CHARGE, FUSION_ALPHA_PARTICLE_ENERGY,ONE_EV from essos.dynamics import Tracing, Particles +from essos.objective_functions import normB_axis # Input parameters tmax = 1.e-4 @@ -22,16 +25,25 @@ rtol=1.e-7 energy=FUSION_ALPHA_PARTICLE_ENERGY - - -# Load coils and field -json_file = os.path.join(os.path.dirname(__name__), 'input_files', 'QH_simple_scaled.json') -coils = Coils_from_simsopt(json_file,nfp=4) +# # Load coils and field +json_file = os.path.join(os.path.dirname(__file__), '..', 'input_files', 'QH_simple_scaled.json') +coils = Coils.from_simsopt(json_file, nfp=4) field = BiotSavart(coils) +# json_file = os.path.join(os.path.dirname(__file__), '..', 'input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') +# coils = Coils.from_json(json_file) +# field = BiotSavart(coils) - +#renormalize coisl to have B_target=5.7 on axis +B_axis_old=normB_axis(field,npoints=200) +#print(jnp.average(B_axis_old)) +B_target=5.7 +coils.dofs_currents=coils.dofs_currents*B_target/jnp.average(B_axis_old) +field=BiotSavart(coils) +#B_axis_new=normB_axis(field,npoints=200) +#print(jnp.average(B_axis_new)) # Load coils and field -wout_file = os.path.join(os.path.dirname(__name__), 'input_files','wout_QH_simple_scaled.nc') +wout_file = os.path.join(os.path.dirname(__file__), '..', 'input_files','wout_QH_simple_scaled.nc') +# wout_file = os.path.join(os.path.dirname(__file__), "../input_files", "wout_LandremanPaul2021_QA_reactorScale_lowres.nc") vmec = Vmec(wout_file) timeI=time() @@ -46,8 +58,8 @@ print(f"Initialization performed") # Trace in ESSOS time0 = time() -tracing = block_until_ready(Tracing(field=field, model='GuidingCenterAdaptative', particles=particles, - maxtime=tmax, timestep=timestep,times_to_trace=times_to_trace, atol=atol,rtol=rtol,boundary=boundary)) +tracing = Tracing(field=field, model='GuidingCenterAdaptative', particles=particles, + maxtime=tmax, timestep=timestep,times_to_trace=times_to_trace, atol=atol,rtol=rtol,boundary=boundary) print(f"ESSOS tracing took {time()-time0:.2f} seconds") print(f"Final loss fraction: {tracing.loss_fractions[-1]*100:.2f}%") trajectories = tracing.trajectories @@ -64,7 +76,7 @@ tracing.plot(ax=ax1, show=False, n_trajectories_plot=nparticles) for i, trajectory in enumerate(trajectories): - ax2.plot(tracing.times, jnp.abs(tracing.energy[i]-particles.energy)/particles.energy, label=f'Particle {i+1}') + ax2.plot(tracing.times, jnp.abs(tracing.energy()[i]-particles.energy)/particles.energy, label=f'Particle {i+1}') ax3.plot(tracing.times, trajectory[:, 3]/particles.total_speed, label=f'Particle {i+1}') #ax4.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') ax4.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') @@ -81,7 +93,6 @@ plt.tight_layout() plt.show() - ## Save results in vtk format to analyze in Paraview # tracing.to_vtk('trajectories') # coils.to_vtk('coils') diff --git a/examples/trace_particles_coils_guidingcenter_with_classifier_scaled_currents.py b/examples/particle_tracing/trace_particles_coils_guidingcenter_with_classifier_scaled_currents.py similarity index 91% rename from examples/trace_particles_coils_guidingcenter_with_classifier_scaled_currents.py rename to examples/particle_tracing/trace_particles_coils_guidingcenter_with_classifier_scaled_currents.py index 364888d1..fd9c1754 100644 --- a/examples/trace_particles_coils_guidingcenter_with_classifier_scaled_currents.py +++ b/examples/particle_tracing/trace_particles_coils_guidingcenter_with_classifier_scaled_currents.py @@ -9,7 +9,7 @@ import matplotlib.pyplot as plt from essos.fields import BiotSavart,Vmec from essos.surfaces import SurfaceClassifier -from essos.coils import Coils_from_json,Coils_from_simsopt +from essos.coils import Coils from essos.constants import ALPHA_PARTICLE_MASS, ALPHA_PARTICLE_CHARGE, FUSION_ALPHA_PARTICLE_ENERGY,ONE_EV from essos.dynamics import Tracing, Particles from essos.objective_functions import normB_axis @@ -25,11 +25,9 @@ rtol=1.e-7 energy=FUSION_ALPHA_PARTICLE_ENERGY - - # Load coils and field -json_file = os.path.join(os.path.dirname(__name__), 'input_files', 'QH_simple_scaled.json') -coils = Coils_from_simsopt(json_file,nfp=4) +json_file = os.path.join(os.path.dirname(__file__), '..', 'input_files', 'QH_simple_scaled.json') +coils = Coils.from_simsopt(json_file,nfp=4) field = BiotSavart(coils) #renormalize coisl to have B_target=5.7 on axis @@ -41,7 +39,7 @@ #B_axis_new=normB_axis(field,npoints=200) #print(jnp.average(B_axis_new)) # Load coils and field -wout_file = os.path.join(os.path.dirname(__name__), 'input_files','wout_QH_simple_scaled.nc') +wout_file = os.path.join(os.path.dirname(__file__), '..', 'input_files','wout_QH_simple_scaled.nc') vmec = Vmec(wout_file) timeI=time() diff --git a/examples/trace_particles_vmec.py b/examples/particle_tracing/trace_particles_vmec.py similarity index 91% rename from examples/trace_particles_vmec.py rename to examples/particle_tracing/trace_particles_vmec.py index abda04b7..cbc3f34e 100644 --- a/examples/trace_particles_vmec.py +++ b/examples/particle_tracing/trace_particles_vmec.py @@ -12,7 +12,7 @@ # Input parameters tmax = 1e-4 timestep = 1.e-8 -times_to_trace=5000 +times_to_trace=1000 nparticles_per_core=6 nparticles = number_of_processors_to_use*nparticles_per_core n_particles_to_plot = 4 @@ -24,7 +24,7 @@ energy=FUSION_ALPHA_PARTICLE_ENERGY # Load coils and field -wout_file = os.path.join(os.path.dirname(__file__), "input_files", "wout_LandremanPaul2021_QA_reactorScale_lowres.nc") +wout_file = os.path.join(os.path.dirname(__file__), "../input_files", "wout_LandremanPaul2021_QA_reactorScale_lowres.nc") vmec = Vmec(wout_file) # Initialize particles @@ -56,7 +56,7 @@ for i in np.random.choice(nparticles, size=n_particles_to_plot, replace=False): trajectory = trajectories[i] ## Plot energy error - ax2.plot(tracing.times[2:], jnp.abs(tracing.energy[i][2:]-particles.energy)/particles.energy, label=f'Particle {i+1}') + ax2.plot(tracing.times[2:], jnp.abs(tracing.energy()[i][2:]-particles.energy)/particles.energy, label=f'Particle {i+1}') ## Plot velocity parallel to the magnetic field ax3.plot(tracing.times, trajectory[:, 3]/particles.total_speed, label=f'Particle {i+1}') ## Plot s-coordinate diff --git a/examples/trace_particles_vmec_Electric_field.py b/examples/particle_tracing/trace_particles_vmec_Electric_field.py similarity index 90% rename from examples/trace_particles_vmec_Electric_field.py rename to examples/particle_tracing/trace_particles_vmec_Electric_field.py index 95f3e720..ce9e380b 100644 --- a/examples/trace_particles_vmec_Electric_field.py +++ b/examples/particle_tracing/trace_particles_vmec_Electric_field.py @@ -24,11 +24,11 @@ energy=FUSION_ALPHA_PARTICLE_ENERGY # Load coils and field -wout_file = os.path.join(os.path.dirname(__file__), "input_files", "wout_LandremanPaul2021_QA_reactorScale_lowres.nc") +wout_file = os.path.join(os.path.dirname(__file__), "..", "input_files", "wout_LandremanPaul2021_QA_reactorScale_lowres.nc") vmec = Vmec(wout_file) #Load electric field -Er_file=os.path.join(os.path.dirname(__file__), 'input_files','Er.h5') +Er_file = os.path.join(os.path.dirname(__file__), "..", "input_files", "Er.h5") Electric_field=Electric_field_flux(Er_filename=Er_file,vmec=vmec) # Initialize particles @@ -61,7 +61,7 @@ for i in np.random.choice(nparticles, size=min(2, nparticles), replace=False): trajectory = trajectories[i] ## Plot energy error - ax2.plot(tracing.times[2:], jnp.abs(tracing.energy[i][2:]-particles.energy)/particles.energy, label=f'Particle {i+1}') + ax2.plot(tracing.times[2:], jnp.abs(tracing.energy()[i][2:]-particles.energy)/particles.energy, label=f'Particle {i+1}') ## Plot velocity parallel to the magnetic field ax3.plot(tracing.times, trajectory[:, 3]/particles.total_speed, label=f'Particle {i+1}') ## Plot s-coordinate diff --git a/examples/particle_tracing/trace_particles_vmec_classifier.py b/examples/particle_tracing/trace_particles_vmec_classifier.py new file mode 100644 index 00000000..10c2e9b5 --- /dev/null +++ b/examples/particle_tracing/trace_particles_vmec_classifier.py @@ -0,0 +1,82 @@ +import os +number_of_processors_to_use = 1 # Parallelization, this should divide nparticles +os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' +from time import time +import jax.numpy as jnp +import matplotlib.pyplot as plt +from essos.fields import Vmec +from essos.constants import ALPHA_PARTICLE_MASS, ALPHA_PARTICLE_CHARGE, FUSION_ALPHA_PARTICLE_ENERGY +from essos.dynamics import Tracing, Particles +from essos.surfaces import SurfaceClassifier +import numpy as np + +# Input parameters +tmax = 1e-4 +timestep = 1.e-8 +times_to_trace=1000 +nparticles_per_core=6 +nparticles = number_of_processors_to_use*nparticles_per_core +n_particles_to_plot = 4 +s = 0.6 # s-coordinate: flux surface label +theta = jnp.linspace(0, 2*jnp.pi, nparticles) +phi = jnp.linspace(0, 2*jnp.pi/2/4, nparticles) +atol = 1e-8 +rtol = 1e-8 +energy=FUSION_ALPHA_PARTICLE_ENERGY + +# Load coils and field +wout_file = os.path.join(os.path.dirname(__file__), "../input_files", "wout_LandremanPaul2021_QA_reactorScale_lowres.nc") +vmec = Vmec(wout_file) +boundary=SurfaceClassifier(vmec.surface,h=0.1) + +# Initialize particles +Z0 = jnp.zeros(nparticles) +phi0 = jnp.zeros(nparticles) +initial_xyz=jnp.array([[s]*nparticles, theta, phi]).T +particles = Particles(initial_xyz=initial_xyz, mass=ALPHA_PARTICLE_MASS, + charge=ALPHA_PARTICLE_CHARGE, energy=energy, field=vmec) +# Trace in ESSOS +time0 = time() +tracing = Tracing(field=vmec, model='GuidingCenterAdaptative', particles=particles, maxtime=tmax, + timestep=timestep,times_to_trace=times_to_trace, atol=atol,rtol=rtol,boundary=boundary) +print(f"ESSOS tracing of {nparticles} particles during {tmax}s took {time()-time0:.2f} seconds") +print(f"Final loss fraction: {tracing.loss_fractions[-1]*100:.2f}%") +trajectories = tracing.trajectories + +# Plot trajectories, velocity parallel to the magnetic field, loss fractions and/or energy error +fig = plt.figure(figsize=(9, 8)) +ax1 = fig.add_subplot(221, projection='3d') +ax2 = fig.add_subplot(222) +ax3 = fig.add_subplot(223) +ax4 = fig.add_subplot(224) + +# Plot 5 random particles +## Plot trajectories in 3D +vmec.surface.plot(ax=ax1, show=False, alpha=0.4) +tracing.plot(ax=ax1, show=False, n_trajectories_plot=nparticles) +for i in np.random.choice(nparticles, size=n_particles_to_plot, replace=False): + trajectory = trajectories[i] + ## Plot energy error + ax2.plot(tracing.times[2:], jnp.abs(tracing.energy()[i][2:]-particles.energy)/particles.energy, label=f'Particle {i+1}') + ## Plot velocity parallel to the magnetic field + ax3.plot(tracing.times, trajectory[:, 3]/particles.total_speed, label=f'Particle {i+1}') + ## Plot s-coordinate + ax4.plot(tracing.times, trajectory[:,0], label=f'Particle {i+1}') + # ax4.set_ylabel(r'$s=\psi/\psi_b$') +## Plot loss fractions +#ax4.plot(tracing.times, tracing.loss_fractions) +#ax4.set_ylabel('Loss Fraction');ax4.set_ylim(0, 1);ax4.set_xscale('log') +ax2.set_xscale('log') +ax2.set_yscale('log') +ax2.set_ylabel('Relative Energy Error') +ax2.set_xlabel('Time (s)') +ax3.set_ylim(-1, 1) +ax3.set_ylabel(r'$v_{\parallel}/v$') +ax3.set_xlabel('Time (s)') +ax4.set_xlabel('Time (s)') +plt.tight_layout() +plt.show() + +# # Save results in vtk format to analyze in Paraview +# vmec.surface.to_vtk('surface') +# tracing.to_vtk('trajectories') diff --git a/examples/particle_tracing_collisions/statistics_collisions_velocity_distributions_mu_Adaptative.py b/examples/particle_tracing_collisions/statistics_collisions_velocity_distributions_mu_Adaptative.py new file mode 100644 index 00000000..70613ea0 --- /dev/null +++ b/examples/particle_tracing_collisions/statistics_collisions_velocity_distributions_mu_Adaptative.py @@ -0,0 +1,336 @@ +import os +number_of_processors_to_use = 1 # Parallelization, this should divide nparticles +os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' +from time import time +import jax.numpy as jnp +import matplotlib.pyplot as plt +import matplotlib.colors +from essos.fields import BiotSavart,Vmec +from essos.constants import PROTON_MASS, ONE_EV,ELECTRON_MASS,SPEED_OF_LIGHT,ELEMENTARY_CHARGE +from essos.dynamics import Tracing, Particles +from essos.background_species import BackgroundSpecies,gamma_ab +import numpy as np +import jax + +from jax import config +# to use higher precision +config.update("jax_enable_x64", True) + +# Input parameters +tmax = 1e-4 +dt=1.e-14 +times_to_trace=1000 +nparticles_per_core=10 +nparticles = number_of_processors_to_use*nparticles_per_core +s=0.25 +num_steps = jnp.round(tmax/dt) +mass=PROTON_MASS +mass_e=ELECTRON_MASS +T_test=3000. +energy=T_test*ONE_EV + +# # Load coils and field +# json_file = os.path.join(os.path.dirname(__file__), '../input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') +# coils = Coils_from_json(json_file) +plt.rcParams.update({'font.size': 16}) +# field = BiotSavart(coils) + +# # Initialize particles +# Z0 = jnp.zeros(nparticles) +# phi0 = jnp.zeros(nparticles) +# initial_xyz=jnp.array([R0*jnp.cos(phi0), R0*jnp.sin(phi0), Z0]).T +# particles = Particles(initial_xyz=initial_xyz,initial_vparallel_over_v=1.0*jnp.ones(nparticles), mass=mass, energy=energy) + + +# Load coils and field +wout_file = os.path.join(os.path.dirname(__file__), '..', 'input_files', "wout_LandremanPaul2021_QA_reactorScale_lowres.nc") +vmec = Vmec(wout_file, ntheta=60, nphi=60, range_torus='half period', close=True) + +theta = jnp.zeros(nparticles) +phi = jnp.zeros(nparticles) + +initial_xyz=jnp.array([[s]*nparticles, theta, phi]).T +particles = Particles(initial_xyz=initial_xyz, mass=mass, + charge=ELEMENTARY_CHARGE, energy=energy, field=vmec,initial_vparallel_over_v=1.0*jnp.ones(nparticles)) + +#Initialize background species +number_species=1 #(electrons,deuterium) +mass_array=jnp.array([1.]) #mass_over_mproton +charge_array=jnp.array([1.]) #mass_over_mproton +T0=1.e+3 #eV +n0=1e+20 #m^-3 +n_array=jnp.array([n0]) +T_array=jnp.array([T0]) +species = BackgroundSpecies(number_species=number_species, mass_array=mass_array, charge_array=charge_array, n_array=n_array, T_array=T_array) +vth_c=jnp.sqrt(T0*ONE_EV/PROTON_MASS)/SPEED_OF_LIGHT +vpar_mean=0. +vpar_sigma=vth_c +v_mean=vth_c*jnp.sqrt(8./jnp.pi) +v_sigma=vth_c*jnp.sqrt((3.*jnp.pi-8.)/jnp.pi) +vperp_mean=vth_c*jnp.sqrt(jnp.pi/2.) +vperp_sigma=vth_c*jnp.sqrt(2.-jnp.pi/2.) +pitch_mean=0. +pitch_sigma=jnp.sqrt(2.**2/12) + + +time0 = time() +tracing = Tracing(field=vmec, model='GuidingCenterCollisionsMuAdaptative', particles=particles, + maxtime=tmax, timestep=dt,times_to_trace=times_to_trace,species=species,tag_gc=0.) +print(f"ESSOS tracing took {time()-time0:.2f} seconds") +trajectories = tracing.trajectories + + +# Plot trajectories, velocity parallel to the magnetic field, and energy error +fig = plt.figure(figsize=(9, 8)) +ax1 = fig.add_subplot(221, projection='3d') +ax2 = fig.add_subplot(222) +ax3 = fig.add_subplot(223) +ax4 = fig.add_subplot(224) + +#vmec.plot(ax=ax1, show=False) +#tracing.plot(ax=ax1, show=False) + +# Plot only a random subset of 10 particles in 3D +subset_size = 10 +import numpy as np +subset_indices = np.random.choice(len(trajectories), subset_size, replace=False) + +for i in subset_indices: + trajectory = trajectories[i] + ax1.plot(trajectory[:,0], trajectory[:,1], trajectory[:,2], label=f'Particle {i+1}') + ax2.plot(tracing.times, (tracing.energy()[i]-tracing.energy()[i,0])/tracing.energy()[i,0], label=f'Particle {i+1}') + ax3.plot(tracing.times, 299792458*trajectory[:, 3]/jnp.sqrt(tracing.energy()[i]/mass*2.), label=f'Particle {i+1}') + ax4.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') + + +# Set bold font for all axes and tick labels +for ax in [ax1, ax2, ax3, ax4]: + ax.xaxis.label.set_fontweight('bold') + ax.yaxis.label.set_fontweight('bold') + ax.title.set_fontweight('bold') + for label in (ax.get_xticklabels() + ax.get_yticklabels()): + label.set_fontweight('bold') + +ax2.set_xlabel(r'$t~[\mathrm{s}]$', fontweight='bold') +ax2.set_ylabel(r'$\frac{E-E_0}{E_0}$', fontweight='bold') +ax3.set_ylabel(r'$v_{\parallel}/v$', fontweight='bold') +ax3.set_xlabel(r'$t~[\mathrm{s}]$', fontweight='bold') +ax4.set_xlabel(r'$R~[\mathrm{m}]$', fontweight='bold') +ax4.set_ylabel(r'$Z~[\mathrm{m}]$', fontweight='bold') +plt.tight_layout() +plt.savefig('traj.pdf') + + +v=jnp.sqrt(tracing.energy()*2./particles.mass) +vpar=trajectories[:,:,3]*SPEED_OF_LIGHT +vpar=jnp.where(jnp.isfinite(vpar), vpar, jnp.nan) +vperp=tracing.v_perp() +pitch=vpar/v + + +# Improve font size for all plots +plt.rcParams.update({'font.size': 18, 'font.weight': 'bold'}) + +# 1. v +fig_v = plt.figure(figsize=(7, 5)) +ax_v_mean = fig_v.add_subplot(211) +ax_v_std = fig_v.add_subplot(212) +for ax in [ax_v_mean, ax_v_std]: + for label in (ax.get_xticklabels() + ax.get_yticklabels()): + label.set_fontweight('bold') +ax_v_mean.plot(tracing.times, jnp.nanmean(v/SPEED_OF_LIGHT, axis=0), linewidth=5) +ax_v_mean.axhline(y=v_mean, color='r', linestyle='--', linewidth=5) +ax_v_mean.set_title(r'$\langle v \rangle$', fontweight='bold') +ax_v_mean.set_xlabel('time', fontweight='bold') +ax_v_mean.set_ylabel(r'$v/c$', fontweight='bold') +ax_v_std.plot(tracing.times, jnp.nanstd(v/SPEED_OF_LIGHT, axis=0), linewidth=5) +ax_v_std.axhline(y=v_sigma, color='r', linestyle='--', linewidth=5) +ax_v_std.set_title(r'$\sigma(v)$', fontweight='bold') +ax_v_std.set_xlabel('time', fontweight='bold') +ax_v_std.set_ylabel(r'$v/c$', fontweight='bold') +plt.tight_layout() +fig_v.savefig('statistics_v.pdf', dpi=300) + +# 2. pitch +fig_pitch = plt.figure(figsize=(7, 5)) +ax_pitch_mean = fig_pitch.add_subplot(211) +ax_pitch_std = fig_pitch.add_subplot(212) +for ax in [ax_pitch_mean, ax_pitch_std]: + for label in (ax.get_xticklabels() + ax.get_yticklabels()): + label.set_fontweight('bold') +ax_pitch_mean.plot(tracing.times, jnp.nanmean(pitch, axis=0), linewidth=5) +ax_pitch_mean.axhline(y=pitch_mean, color='r', linestyle='--', linewidth=5) +ax_pitch_mean.set_title(r'$\langle \text{pitch} \rangle$', fontweight='bold') +ax_pitch_mean.set_xlabel('time', fontweight='bold') +ax_pitch_mean.set_ylabel('pitch', fontweight='bold') +ax_pitch_std.plot(tracing.times, jnp.nanstd(pitch, axis=0), linewidth=5) +ax_pitch_std.axhline(y=pitch_sigma, color='r', linestyle='--', linewidth=5) +ax_pitch_std.set_title(r'$\sigma(\text{pitch})$', fontweight='bold') +ax_pitch_std.set_xlabel('time', fontweight='bold') +ax_pitch_std.set_ylabel('pitch', fontweight='bold') +plt.tight_layout() +fig_pitch.savefig('statistics_pitch.pdf', dpi=300) + +# 3. v_parallel/c +fig_vpar = plt.figure(figsize=(7, 5)) +ax_vpar_mean = fig_vpar.add_subplot(211) +ax_vpar_std = fig_vpar.add_subplot(212) +for ax in [ax_vpar_mean, ax_vpar_std]: + for label in (ax.get_xticklabels() + ax.get_yticklabels()): + label.set_fontweight('bold') +ax_vpar_mean.plot(tracing.times, jnp.nanmean(vpar/SPEED_OF_LIGHT, axis=0), linewidth=5) +ax_vpar_mean.axhline(y=vpar_mean, color='r', linestyle='--', linewidth=5) +ax_vpar_mean.set_title(r'$\langle v_{\parallel}/c \rangle$', fontweight='bold') +ax_vpar_mean.set_xlabel(r'$t~[\mathrm{s}]$', fontweight='bold') +ax_vpar_mean.set_ylabel(r'$v_{\parallel}/c$', fontweight='bold') +ax_vpar_std.plot(tracing.times, jnp.nanstd(vpar/SPEED_OF_LIGHT, axis=0), linewidth=5) +ax_vpar_std.axhline(y=vpar_sigma, color='r', linestyle='--', linewidth=5) +ax_vpar_std.set_title(r'$\sigma(v_{\parallel}/c)$', fontweight='bold') +ax_vpar_std.set_xlabel(r'$t~[\mathrm{s}]$', fontweight='bold') +ax_vpar_std.set_ylabel(r'$\sigma_{v_{\parallel}/c}$', fontweight='bold') +plt.tight_layout() +fig_vpar.savefig('statistics_vpar.pdf', dpi=300) + +# 4. v_perp/c +fig_vperp = plt.figure(figsize=(7, 5)) +ax_vperp_mean = fig_vperp.add_subplot(211) +ax_vperp_std = fig_vperp.add_subplot(212) +for ax in [ax_vperp_mean, ax_vperp_std]: + for label in (ax.get_xticklabels() + ax.get_yticklabels()): + label.set_fontweight('bold') +ax_vperp_mean.plot(tracing.times, jnp.nanmean(vperp/SPEED_OF_LIGHT, axis=0), linewidth=5) +ax_vperp_mean.axhline(y=vperp_mean, color='r', linestyle='--', linewidth=5) +ax_vperp_mean.set_title(r'$\langle v_{\perp}/c \rangle$', fontweight='bold') +ax_vperp_mean.set_xlabel(r'$t~[\mathrm{s}]$', fontweight='bold') +ax_vperp_mean.set_ylabel(r'$v_{\perp}/c$', fontweight='bold') +ax_vperp_std.plot(tracing.times, jnp.nanstd(vperp/SPEED_OF_LIGHT, axis=0), linewidth=5) +ax_vperp_std.axhline(y=vperp_sigma, color='r', linestyle='--', linewidth=5) +ax_vperp_std.set_title(r'$\sigma(v_{\perp}/c)$', fontweight='bold') +ax_vperp_std.set_xlabel(r'$t~[\mathrm{s}]$', fontweight='bold') +ax_vperp_std.set_ylabel(r'$\sigma_{v_{\perp}/c}$', fontweight='bold') +plt.tight_layout() +fig_vperp.savefig('statistics_vperp.pdf', dpi=300) + + +# Plot distribution in velocities initial t and final +fig2 = plt.figure(figsize=(9, 8)) +ax12 = fig2.add_subplot(251) +ax22 = fig2.add_subplot(252) +ax32 = fig2.add_subplot(253) +ax42 = fig2.add_subplot(254) +ax52 = fig2.add_subplot(255) +ax62 = fig2.add_subplot(256) +ax72 = fig2.add_subplot(257) +ax82 = fig2.add_subplot(258) +nbins=64 + +v0=jnp.sqrt(tracing.energy()[:,0]*2./particles.mass)/SPEED_OF_LIGHT +vfinal=jnp.sqrt(tracing.energy()[:,-1]*2./particles.mass)/SPEED_OF_LIGHT +vperp0=tracing.v_perp()[:,0]/SPEED_OF_LIGHT +vperpfinal=tracing.v_perp()[:,-1]/SPEED_OF_LIGHT +vpar0=vpar[:,0]/SPEED_OF_LIGHT +vparfinal=vpar[:,-1]/SPEED_OF_LIGHT +pitch0=vpar0/v0 +pitch_final=vparfinal/vfinal + + +bad_indices_v0 = jnp.isnan(v0) +bad_indices_vfinal = jnp.isnan(vfinal) +bad_indices_pitch0 = jnp.isnan(pitch0) +bad_indices_pitch_final = jnp.isnan(pitch_final) +bad_indices_vperp0 = jnp.isnan(vperp0) +bad_indices_vperp_final = jnp.isnan(vperpfinal) +bad_indices_vpar0 = jnp.isnan(vpar0) +bad_indices_vpar_final = jnp.isnan(vparfinal) +good_indices_v0 = ~bad_indices_v0 +good_indices_vfinal = ~bad_indices_vfinal +good_indices_pitch0 = ~bad_indices_pitch0 +good_indices_pitch_final = ~bad_indices_pitch_final +good_indices_vpar0 = ~bad_indices_vpar0 +good_indices_vpar_final = ~bad_indices_vpar_final +good_indices_vperp0 = ~bad_indices_vperp0 +good_indices_vperp_final = ~bad_indices_vperp_final +good_v0 = v0[good_indices_v0] +good_vfinal = vfinal[good_indices_vfinal] +good_pitch0 = pitch0[good_indices_pitch0] +good_pitch_final = pitch_final[good_indices_pitch_final] + + +good_vpar0 = vpar0[good_indices_vpar0] +good_vpar_final = vparfinal[good_indices_vpar_final] +good_vperp0 = vperp0[good_indices_vperp0] +good_vperp_final = vperpfinal[good_indices_vperp_final] + + +v0_counts,v0_bins=jnp.histogram(good_v0,bins=nbins) +vfinal_counts,vfinal_bins=jnp.histogram(good_vfinal,bins=nbins) + +pitch_t0_counts,pitch_t0_bins=jnp.histogram(good_pitch0,bins=nbins) +pitch_tfinal_counts,pitch_tfinal_bins=jnp.histogram(good_pitch_final,bins=nbins) + +vpar_t0_counts,vpar_t0_bins=jnp.histogram(good_vpar0,bins=nbins) +vpar_tfinal_counts,vpar_tfinal_bins=jnp.histogram(good_vpar_final,bins=nbins) + + +vperp_t0_counts,vperp_t0_bins=jnp.histogram(good_vperp0,bins=nbins) +vperp_tfinal_counts,vperp_tfinal_bins=jnp.histogram(good_vperp_final,bins=nbins) + + +ax12.stairs(v0_counts,v0_bins) +ax22.stairs(vfinal_counts,vfinal_bins) +ax32.stairs(vpar_t0_counts,vpar_t0_bins) +ax42.stairs(vpar_tfinal_counts,vpar_tfinal_bins) +ax52.stairs(pitch_t0_counts,pitch_t0_bins) +ax62.stairs(pitch_tfinal_counts,pitch_tfinal_bins) +ax72.stairs(vperp_t0_counts,vperp_t0_bins) +ax82.stairs(vperp_tfinal_counts,vperp_tfinal_bins) +plt.figure(figsize=(7, 5)) +plt.hist(good_vfinal, bins=nbins, color='b', edgecolor='black', alpha=0.7) +plt.axvline(np.mean(good_v0), color='r', linestyle='--', linewidth=3, label='Initial Mean') +plt.title(r'$v/c$ Distribution', fontweight='bold') +plt.xlabel(r'$v/c$', fontweight='bold') +plt.ylabel('Counts', fontweight='bold') +plt.legend(fontsize=14) +plt.tight_layout() +plt.savefig('dist_v.pdf', dpi=300) + +plt.figure(figsize=(7, 5)) +plt.hist(good_pitch_final, bins=nbins, color='g', edgecolor='black', alpha=0.7) +plt.axvline(np.mean(good_pitch0), color='r', linestyle='--', linewidth=3, label='Initial Mean') +plt.title(r'Pitch Distribution', fontweight='bold') +plt.xlabel(r'Pitch', fontweight='bold') +plt.ylabel('Counts', fontweight='bold') +plt.legend(fontsize=14) +plt.tight_layout() +plt.savefig('dist_pitch.pdf', dpi=300) + +plt.figure(figsize=(7, 5)) +plt.hist(good_vpar_final, bins=nbins, color='#FA7000', edgecolor='black', alpha=0.7) +plt.axvline(np.mean(good_vpar0), color='b', linestyle='--', linewidth=3, label='Initial Mean') +plt.title(r'$v_{\parallel}/c$ Distribution', fontweight='bold') +plt.xlabel(r'$v_{\parallel}/c$', fontweight='bold') +plt.ylabel('Counts', fontweight='bold') +plt.legend(fontsize=14) +plt.tight_layout() +plt.savefig('dist_vpar.pdf', dpi=300) + +plt.figure(figsize=(7, 5)) +plt.hist(good_vperp_final, bins=nbins, color='m', edgecolor='black', alpha=0.7) +plt.axvline(np.mean(good_vperp0), color='b', linestyle='--', linewidth=3, label='Initial Mean') +plt.title(r'$v_{\perp}/c$ Distribution', fontweight='bold') +plt.xlabel(r'$v_{\perp}/c$', fontweight='bold') +plt.ylabel('Counts', fontweight='bold') +plt.legend(fontsize=14) +plt.tight_layout() +plt.savefig('dist_vperp.pdf', dpi=300) + + +plt.figure(figsize=(7, 5)) +plt.hist(good_vperp_final, bins=nbins, color='#FA7000', edgecolor='black', alpha=0.7) +plt.axvline(np.mean(good_vperp0), color='b', linestyle='--', linewidth=3, label='Initial Mean') +plt.title(r'$v_{\perp}/c$ Distribution', fontweight='bold') +plt.xlabel(r'$v_{\perp}/c$', fontweight='bold') +plt.ylabel('Counts', fontweight='bold') +plt.legend(fontsize=14) +plt.tight_layout() +plt.savefig('dist_vperp_color.pdf', dpi=300) diff --git a/examples/particle_tracing_collisions/statistics_collisions_velocity_distributions_mu_Fixed.py b/examples/particle_tracing_collisions/statistics_collisions_velocity_distributions_mu_Fixed.py new file mode 100644 index 00000000..41d3e038 --- /dev/null +++ b/examples/particle_tracing_collisions/statistics_collisions_velocity_distributions_mu_Fixed.py @@ -0,0 +1,337 @@ +import os +number_of_processors_to_use = 1 # Parallelization, this should divide nparticles +os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' +from time import time +import jax.numpy as jnp +import matplotlib.pyplot as plt +import matplotlib.colors +from essos.fields import BiotSavart,Vmec +from essos.constants import PROTON_MASS, ONE_EV,ELECTRON_MASS,SPEED_OF_LIGHT,ELEMENTARY_CHARGE +from essos.dynamics import Tracing, Particles +from essos.background_species import BackgroundSpecies,gamma_ab +import numpy as np +import jax + +from jax import config +# to use higher precision +config.update("jax_enable_x64", True) + +# Input parameters +tmax = 1.e-4 +dt=1.e-8 +times_to_trace=1000 +nparticles_per_core=10 +nparticles = number_of_processors_to_use*nparticles_per_core +s=0.25 +num_steps = jnp.round(tmax/dt) +mass=PROTON_MASS +mass_e=ELECTRON_MASS +T_test=3000. +energy=T_test*ONE_EV + + +# # Load coils and field +# json_file = os.path.join(os.path.dirname(__file__), '../input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') +# coils = Coils_from_json(json_file) +plt.rcParams.update({'font.size': 16}) +# field = BiotSavart(coils) + +# # Initialize particles +# Z0 = jnp.zeros(nparticles) +# phi0 = jnp.zeros(nparticles) +# initial_xyz=jnp.array([R0*jnp.cos(phi0), R0*jnp.sin(phi0), Z0]).T +# particles = Particles(initial_xyz=initial_xyz,initial_vparallel_over_v=1.0*jnp.ones(nparticles), mass=mass, energy=energy) + + +# Load coils and field +wout_file = os.path.join(os.path.dirname(__file__), '..', 'input_files', "wout_LandremanPaul2021_QA_reactorScale_lowres.nc") +vmec = Vmec(wout_file, ntheta=60, nphi=60, range_torus='half period', close=True) + +theta = jnp.zeros(nparticles) +phi = jnp.zeros(nparticles) + +initial_xyz=jnp.array([[s]*nparticles, theta, phi]).T +particles = Particles(initial_xyz=initial_xyz, mass=mass, + charge=ELEMENTARY_CHARGE, energy=energy, field=vmec,initial_vparallel_over_v=1.0*jnp.ones(nparticles)) + +#Initialize background species +number_species=1 #(electrons,deuterium) +mass_array=jnp.array([1.]) #mass_over_mproton +charge_array=jnp.array([1.]) #mass_over_mproton +T0=1.e+3 #eV +n0=1e+20 #m^-3 +n_array=jnp.array([n0]) +T_array=jnp.array([T0]) +species = BackgroundSpecies(number_species=number_species, mass_array=mass_array, charge_array=charge_array, n_array=n_array, T_array=T_array) +vth_c=jnp.sqrt(T0*ONE_EV/PROTON_MASS)/SPEED_OF_LIGHT +vpar_mean=0. +vpar_sigma=vth_c +v_mean=vth_c*jnp.sqrt(8./jnp.pi) +v_sigma=vth_c*jnp.sqrt((3.*jnp.pi-8.)/jnp.pi) +vperp_mean=vth_c*jnp.sqrt(jnp.pi/2.) +vperp_sigma=vth_c*jnp.sqrt(2.-jnp.pi/2.) +pitch_mean=0. +pitch_sigma=jnp.sqrt(2.**2/12) + + +time0 = time() +tracing = Tracing(field=vmec, model='GuidingCenterCollisionsMuFixed', particles=particles, + maxtime=tmax, timestep=dt,times_to_trace=times_to_trace,species=species,tag_gc=0.) +print(f"ESSOS tracing took {time()-time0:.2f} seconds") +trajectories = tracing.trajectories + + +# Plot trajectories, velocity parallel to the magnetic field, and energy error +fig = plt.figure(figsize=(9, 8)) +ax1 = fig.add_subplot(221, projection='3d') +ax2 = fig.add_subplot(222) +ax3 = fig.add_subplot(223) +ax4 = fig.add_subplot(224) + +#vmec.plot(ax=ax1, show=False) +#tracing.plot(ax=ax1, show=False) + +# Plot only a random subset of 10 particles in 3D +subset_size = 10 +import numpy as np +subset_indices = np.random.choice(len(trajectories), subset_size, replace=False) + +for i in subset_indices: + trajectory = trajectories[i] + ax1.plot(trajectory[:,0], trajectory[:,1], trajectory[:,2], label=f'Particle {i+1}') + ax2.plot(tracing.times, (tracing.energy()[i]-tracing.energy()[i,0])/tracing.energy()[i,0], label=f'Particle {i+1}') + ax3.plot(tracing.times, 299792458*trajectory[:, 3]/jnp.sqrt(tracing.energy()[i]/mass*2.), label=f'Particle {i+1}') + ax4.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') + + +# Set bold font for all axes and tick labels +for ax in [ax1, ax2, ax3, ax4]: + ax.xaxis.label.set_fontweight('bold') + ax.yaxis.label.set_fontweight('bold') + ax.title.set_fontweight('bold') + for label in (ax.get_xticklabels() + ax.get_yticklabels()): + label.set_fontweight('bold') + +ax2.set_xlabel(r'$t~[\mathrm{s}]$', fontweight='bold') +ax2.set_ylabel(r'$\frac{E-E_0}{E_0}$', fontweight='bold') +ax3.set_ylabel(r'$v_{\parallel}/v$', fontweight='bold') +ax3.set_xlabel(r'$t~[\mathrm{s}]$', fontweight='bold') +ax4.set_xlabel(r'$R~[\mathrm{m}]$', fontweight='bold') +ax4.set_ylabel(r'$Z~[\mathrm{m}]$', fontweight='bold') +plt.tight_layout() +plt.savefig('traj.pdf') + +v=jnp.sqrt(tracing.energy()*2./particles.mass) +vpar=trajectories[:,:,3]*SPEED_OF_LIGHT +vpar=jnp.where(jnp.isfinite(vpar), vpar, jnp.nan) +vperp=tracing.v_perp() +pitch=vpar/v + + +# Improve font size for all plots +plt.rcParams.update({'font.size': 18, 'font.weight': 'bold'}) + +# 1. v +fig_v = plt.figure(figsize=(7, 5)) +ax_v_mean = fig_v.add_subplot(211) +ax_v_std = fig_v.add_subplot(212) +for ax in [ax_v_mean, ax_v_std]: + for label in (ax.get_xticklabels() + ax.get_yticklabels()): + label.set_fontweight('bold') +ax_v_mean.plot(tracing.times, jnp.nanmean(v/SPEED_OF_LIGHT, axis=0), linewidth=5) +ax_v_mean.axhline(y=v_mean, color='r', linestyle='--', linewidth=5) +ax_v_mean.set_title(r'$\langle v \rangle$', fontweight='bold') +ax_v_mean.set_xlabel('time', fontweight='bold') +ax_v_mean.set_ylabel(r'$v/c$', fontweight='bold') +ax_v_std.plot(tracing.times, jnp.nanstd(v/SPEED_OF_LIGHT, axis=0), linewidth=5) +ax_v_std.axhline(y=v_sigma, color='r', linestyle='--', linewidth=5) +ax_v_std.set_title(r'$\sigma(v)$', fontweight='bold') +ax_v_std.set_xlabel('time', fontweight='bold') +ax_v_std.set_ylabel(r'$v/c$', fontweight='bold') +plt.tight_layout() +fig_v.savefig('statistics_v.pdf', dpi=300) + +# 2. pitch +fig_pitch = plt.figure(figsize=(7, 5)) +ax_pitch_mean = fig_pitch.add_subplot(211) +ax_pitch_std = fig_pitch.add_subplot(212) +for ax in [ax_pitch_mean, ax_pitch_std]: + for label in (ax.get_xticklabels() + ax.get_yticklabels()): + label.set_fontweight('bold') +ax_pitch_mean.plot(tracing.times, jnp.nanmean(pitch, axis=0), linewidth=5) +ax_pitch_mean.axhline(y=pitch_mean, color='r', linestyle='--', linewidth=5) +ax_pitch_mean.set_title(r'$\langle \text{pitch} \rangle$', fontweight='bold') +ax_pitch_mean.set_xlabel('time', fontweight='bold') +ax_pitch_mean.set_ylabel('pitch', fontweight='bold') +ax_pitch_std.plot(tracing.times, jnp.nanstd(pitch, axis=0), linewidth=5) +ax_pitch_std.axhline(y=pitch_sigma, color='r', linestyle='--', linewidth=5) +ax_pitch_std.set_title(r'$\sigma(\text{pitch})$', fontweight='bold') +ax_pitch_std.set_xlabel('time', fontweight='bold') +ax_pitch_std.set_ylabel('pitch', fontweight='bold') +plt.tight_layout() +fig_pitch.savefig('statistics_pitch.pdf', dpi=300) + +# 3. v_parallel/c +fig_vpar = plt.figure(figsize=(7, 5)) +ax_vpar_mean = fig_vpar.add_subplot(211) +ax_vpar_std = fig_vpar.add_subplot(212) +for ax in [ax_vpar_mean, ax_vpar_std]: + for label in (ax.get_xticklabels() + ax.get_yticklabels()): + label.set_fontweight('bold') +ax_vpar_mean.plot(tracing.times, jnp.nanmean(vpar/SPEED_OF_LIGHT, axis=0), linewidth=5) +ax_vpar_mean.axhline(y=vpar_mean, color='r', linestyle='--', linewidth=5) +ax_vpar_mean.set_title(r'$\langle v_{\parallel}/c \rangle$', fontweight='bold') +ax_vpar_mean.set_xlabel(r'$t~[\mathrm{s}]$', fontweight='bold') +ax_vpar_mean.set_ylabel(r'$v_{\parallel}/c$', fontweight='bold') +ax_vpar_std.plot(tracing.times, jnp.nanstd(vpar/SPEED_OF_LIGHT, axis=0), linewidth=5) +ax_vpar_std.axhline(y=vpar_sigma, color='r', linestyle='--', linewidth=5) +ax_vpar_std.set_title(r'$\sigma(v_{\parallel}/c)$', fontweight='bold') +ax_vpar_std.set_xlabel(r'$t~[\mathrm{s}]$', fontweight='bold') +ax_vpar_std.set_ylabel(r'$\sigma_{v_{\parallel}/c}$', fontweight='bold') +plt.tight_layout() +fig_vpar.savefig('statistics_vpar.pdf', dpi=300) + +# 4. v_perp/c +fig_vperp = plt.figure(figsize=(7, 5)) +ax_vperp_mean = fig_vperp.add_subplot(211) +ax_vperp_std = fig_vperp.add_subplot(212) +for ax in [ax_vperp_mean, ax_vperp_std]: + for label in (ax.get_xticklabels() + ax.get_yticklabels()): + label.set_fontweight('bold') +ax_vperp_mean.plot(tracing.times, jnp.nanmean(vperp/SPEED_OF_LIGHT, axis=0), linewidth=5) +ax_vperp_mean.axhline(y=vperp_mean, color='r', linestyle='--', linewidth=5) +ax_vperp_mean.set_title(r'$\langle v_{\perp}/c \rangle$', fontweight='bold') +ax_vperp_mean.set_xlabel(r'$t~[\mathrm{s}]$', fontweight='bold') +ax_vperp_mean.set_ylabel(r'$v_{\perp}/c$', fontweight='bold') +ax_vperp_std.plot(tracing.times, jnp.nanstd(vperp/SPEED_OF_LIGHT, axis=0), linewidth=5) +ax_vperp_std.axhline(y=vperp_sigma, color='r', linestyle='--', linewidth=5) +ax_vperp_std.set_title(r'$\sigma(v_{\perp}/c)$', fontweight='bold') +ax_vperp_std.set_xlabel(r'$t~[\mathrm{s}]$', fontweight='bold') +ax_vperp_std.set_ylabel(r'$\sigma_{v_{\perp}/c}$', fontweight='bold') +plt.tight_layout() +fig_vperp.savefig('statistics_vperp.pdf', dpi=300) + + + +# Plot distribution in velocities initial t and final +fig2 = plt.figure(figsize=(9, 8)) +ax12 = fig2.add_subplot(251) +ax22 = fig2.add_subplot(252) +ax32 = fig2.add_subplot(253) +ax42 = fig2.add_subplot(254) +ax52 = fig2.add_subplot(255) +ax62 = fig2.add_subplot(256) +ax72 = fig2.add_subplot(257) +ax82 = fig2.add_subplot(258) +nbins=64 + +v0=jnp.sqrt(tracing.energy()[:,0]*2./particles.mass)/SPEED_OF_LIGHT +vfinal=jnp.sqrt(tracing.energy()[:,-1]*2./particles.mass)/SPEED_OF_LIGHT +vperp0=tracing.v_perp()[:,0]/SPEED_OF_LIGHT +vperpfinal=tracing.v_perp()[:,-1]/SPEED_OF_LIGHT +vpar0=vpar[:,0]/SPEED_OF_LIGHT +vparfinal=vpar[:,-1]/SPEED_OF_LIGHT +pitch0=vpar0/v0 +pitch_final=vparfinal/vfinal + + +bad_indices_v0 = jnp.isnan(v0) +bad_indices_vfinal = jnp.isnan(vfinal) +bad_indices_pitch0 = jnp.isnan(pitch0) +bad_indices_pitch_final = jnp.isnan(pitch_final) +bad_indices_vperp0 = jnp.isnan(vperp0) +bad_indices_vperp_final = jnp.isnan(vperpfinal) +bad_indices_vpar0 = jnp.isnan(vpar0) +bad_indices_vpar_final = jnp.isnan(vparfinal) +good_indices_v0 = ~bad_indices_v0 +good_indices_vfinal = ~bad_indices_vfinal +good_indices_pitch0 = ~bad_indices_pitch0 +good_indices_pitch_final = ~bad_indices_pitch_final +good_indices_vpar0 = ~bad_indices_vpar0 +good_indices_vpar_final = ~bad_indices_vpar_final +good_indices_vperp0 = ~bad_indices_vperp0 +good_indices_vperp_final = ~bad_indices_vperp_final +good_v0 = v0[good_indices_v0] +good_vfinal = vfinal[good_indices_vfinal] +good_pitch0 = pitch0[good_indices_pitch0] +good_pitch_final = pitch_final[good_indices_pitch_final] + + +good_vpar0 = vpar0[good_indices_vpar0] +good_vpar_final = vparfinal[good_indices_vpar_final] +good_vperp0 = vperp0[good_indices_vperp0] +good_vperp_final = vperpfinal[good_indices_vperp_final] + + +v0_counts,v0_bins=jnp.histogram(good_v0,bins=nbins) +vfinal_counts,vfinal_bins=jnp.histogram(good_vfinal,bins=nbins) + +pitch_t0_counts,pitch_t0_bins=jnp.histogram(good_pitch0,bins=nbins) +pitch_tfinal_counts,pitch_tfinal_bins=jnp.histogram(good_pitch_final,bins=nbins) + +vpar_t0_counts,vpar_t0_bins=jnp.histogram(good_vpar0,bins=nbins) +vpar_tfinal_counts,vpar_tfinal_bins=jnp.histogram(good_vpar_final,bins=nbins) + + +vperp_t0_counts,vperp_t0_bins=jnp.histogram(good_vperp0,bins=nbins) +vperp_tfinal_counts,vperp_tfinal_bins=jnp.histogram(good_vperp_final,bins=nbins) + + +ax12.stairs(v0_counts,v0_bins) +ax22.stairs(vfinal_counts,vfinal_bins) +ax32.stairs(vpar_t0_counts,vpar_t0_bins) +ax42.stairs(vpar_tfinal_counts,vpar_tfinal_bins) +ax52.stairs(pitch_t0_counts,pitch_t0_bins) +ax62.stairs(pitch_tfinal_counts,pitch_tfinal_bins) +ax72.stairs(vperp_t0_counts,vperp_t0_bins) +ax82.stairs(vperp_tfinal_counts,vperp_tfinal_bins) +plt.figure(figsize=(7, 5)) +plt.hist(good_vfinal, bins=nbins, color='b', edgecolor='black', alpha=0.7) +plt.axvline(np.mean(good_v0), color='r', linestyle='--', linewidth=3, label='Initial Mean') +plt.title(r'$v/c$ Distribution', fontweight='bold') +plt.xlabel(r'$v/c$', fontweight='bold') +plt.ylabel('Counts', fontweight='bold') +plt.legend(fontsize=14) +plt.tight_layout() +plt.savefig('dist_v.pdf', dpi=300) + +plt.figure(figsize=(7, 5)) +plt.hist(good_pitch_final, bins=nbins, color='g', edgecolor='black', alpha=0.7) +plt.axvline(np.mean(good_pitch0), color='r', linestyle='--', linewidth=3, label='Initial Mean') +plt.title(r'Pitch Distribution', fontweight='bold') +plt.xlabel(r'Pitch', fontweight='bold') +plt.ylabel('Counts', fontweight='bold') +plt.legend(fontsize=14) +plt.tight_layout() +plt.savefig('dist_pitch.pdf', dpi=300) + +plt.figure(figsize=(7, 5)) +plt.hist(good_vpar_final, bins=nbins, color='#FA7000', edgecolor='black', alpha=0.7) +plt.axvline(np.mean(good_vpar0), color='b', linestyle='--', linewidth=3, label='Initial Mean') +plt.title(r'$v_{\parallel}/c$ Distribution', fontweight='bold') +plt.xlabel(r'$v_{\parallel}/c$', fontweight='bold') +plt.ylabel('Counts', fontweight='bold') +plt.legend(fontsize=14) +plt.tight_layout() +plt.savefig('dist_vpar.pdf', dpi=300) + +plt.figure(figsize=(7, 5)) +plt.hist(good_vperp_final, bins=nbins, color='m', edgecolor='black', alpha=0.7) +plt.axvline(np.mean(good_vperp0), color='b', linestyle='--', linewidth=3, label='Initial Mean') +plt.title(r'$v_{\perp}/c$ Distribution', fontweight='bold') +plt.xlabel(r'$v_{\perp}/c$', fontweight='bold') +plt.ylabel('Counts', fontweight='bold') +plt.legend(fontsize=14) +plt.tight_layout() +plt.savefig('dist_vperp.pdf', dpi=300) + + +plt.figure(figsize=(7, 5)) +plt.hist(good_vperp_final, bins=nbins, color='#FA7000', edgecolor='black', alpha=0.7) +plt.axvline(np.mean(good_vperp0), color='b', linestyle='--', linewidth=3, label='Initial Mean') +plt.title(r'$v_{\perp}/c$ Distribution', fontweight='bold') +plt.xlabel(r'$v_{\perp}/c$', fontweight='bold') +plt.ylabel('Counts', fontweight='bold') +plt.legend(fontsize=14) +plt.tight_layout() +plt.savefig('dist_vperp_color.pdf', dpi=300) diff --git a/examples/testing_collisions_velocity_distributions_mu_time.py b/examples/particle_tracing_collisions/statistics_collisions_velocity_distributions_mu_time.py similarity index 76% rename from examples/testing_collisions_velocity_distributions_mu_time.py rename to examples/particle_tracing_collisions/statistics_collisions_velocity_distributions_mu_time.py index f3dc4449..ec10d276 100644 --- a/examples/testing_collisions_velocity_distributions_mu_time.py +++ b/examples/particle_tracing_collisions/statistics_collisions_velocity_distributions_mu_time.py @@ -5,9 +5,8 @@ import jax.numpy as jnp import matplotlib.pyplot as plt import matplotlib.colors -from essos.fields import BiotSavart -from essos.coils import Coils_from_json -from essos.constants import PROTON_MASS, ONE_EV,ELECTRON_MASS,SPEED_OF_LIGHT +from essos.fields import BiotSavart,Vmec +from essos.constants import PROTON_MASS, ONE_EV,ELECTRON_MASS,SPEED_OF_LIGHT,ELEMENTARY_CHARGE from essos.dynamics import Tracing, Particles from essos.background_species import BackgroundSpecies,gamma_ab import numpy as np @@ -15,13 +14,13 @@ # Input parameters light_speed=SPEED_OF_LIGHT -tmax = 1.e-5 +tmax = 1.e-4 dt=1.e-8 -nparticles_per_core=10 +nparticles_per_core=100 nparticles = number_of_processors_to_use*nparticles_per_core -R0 = 1.25#jnp.linspace(1.23, 1.27, nparticles) +s=0.25 trace_tolerance = 1e-7 -times_to_trace=100 +times_to_trace=1000 mass=PROTON_MASS mass_a=4.*mass mass_e=ELECTRON_MASS @@ -42,15 +41,15 @@ # Load coils and field -json_file = os.path.join(os.path.dirname(__file__), 'input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') -coils = Coils_from_json(json_file) -field = BiotSavart(coils) +wout_file = os.path.join(os.path.dirname(__file__), '..', 'input_files', "wout_LandremanPaul2021_QA_reactorScale_lowres.nc") +vmec = Vmec(wout_file, ntheta=60, nphi=60, range_torus='half period', close=True) -# Initialize particles -Z0 = jnp.zeros(nparticles) -phi0 = jnp.zeros(nparticles) -initial_xyz=jnp.array([R0*jnp.cos(phi0), R0*jnp.sin(phi0), Z0]).T -particles = Particles(initial_xyz=initial_xyz,initial_vparallel_over_v=-1.*jnp.ones(nparticles), mass=mass, energy=energy) +theta = jnp.zeros(nparticles) +phi = jnp.zeros(nparticles) + +initial_xyz=jnp.array([[s]*nparticles, theta, phi]).T +particles = Particles(initial_xyz=initial_xyz, mass=mass, + charge=ELEMENTARY_CHARGE, energy=energy, field=vmec,initial_vparallel_over_v=1.0*jnp.ones(nparticles)) #Initialize background species @@ -82,67 +81,47 @@ pitch_mean=0. pitch_sigma=jnp.sqrt(2.**2/12) -#import jax -#import jax.numpy as jnp -#from essos.dynamics import GuidingCenterCollisionsDriftMu as GCCD -#from essos.dynamics import GuidingCenterCollisionsDiffusionMu as GCCDiff -#from essos.background_species import nu_s_ab,nu_D_ab,nu_par_ab, d_nu_par_ab -#B_particle=jax.vmap(field.AbsB,in_axes=0)(particles.initial_xyz) -#mu=particles.initial_vperpendicular**2*particles.mass*0.5/B_particle/particles.mass -#initial_conditions = jnp.concatenate([particles.initial_xyz,particles.initial_vparallel[:, None],mu[:, None]],axis=1) -#args = (field, particles,species) -#GCCD(0,initial_conditions[0],args) -#GCCDiff(0,initial_conditions[0],args) -#initial_condition=initial_conditions[0] -#initial_condition = jnp.concatenate([particles.initial_xyz,total_speed_temp[:, None], particles.initial_vparallel_over_v[:, None]], axis=1)[0] -#initial_condition = jnp.concatenate([particles.initial_xyz,total_speed_temp[:, None], particles.initial_vparallel_over_v[:, None]], axis=1)[0] - # Trace in ESSOS time0 = time() -tracing = Tracing(field=field, model='GuidingCenterCollisionsMuFixed', particles=particles, +tracing = Tracing(field=vmec, model='GuidingCenterCollisionsMuFixed', particles=particles, maxtime=tmax, timestep=dt,times_to_trace=times_to_trace,species=species,tag_gc=0.) print(f"ESSOS tracing took {time()-time0:.2f} seconds") trajectories = tracing.trajectories - -# Plot trajectories, velocity parallel to the magnetic field, and energy error fig = plt.figure(figsize=(9, 8)) ax1 = fig.add_subplot(221, projection='3d') ax2 = fig.add_subplot(222) ax3 = fig.add_subplot(223) ax4 = fig.add_subplot(224) -coils.plot(ax=ax1, show=False) -tracing.plot(ax=ax1, show=False) +#vmec.plot(ax=ax1, show=False) +#tracing.plot(ax=ax1, show=False) -v=jnp.sqrt(tracing.energy*2./particles.mass) +# Plot only a random subset of 10 particles in 3D +subset_size = 10 +subset_indices = np.random.choice(len(trajectories), subset_size, replace=False) -for i, trajectory in enumerate(trajectories): - #ax2.plot(tracing.times, (tracing.energy[i]-tracing.energy[i,0])/tracing.energy[i,0], label=f'Particle {i+1}') - ax2.plot(tracing.times, (v[i]-v[i,0])/v[i,0], label=f'Particle {i+1}') - ax3.plot(tracing.times, trajectory[:, 3]/jnp.sqrt(tracing.energy[i]/mass*2.), label=f'Particle {i+1}') +for i in subset_indices: + trajectory = trajectories[i] + ax1.plot(trajectory[:,0], trajectory[:,1], trajectory[:,2], label=f'Particle {i+1}') + ax2.plot(tracing.times, (tracing.energy()[i]-tracing.energy()[i,0])/tracing.energy()[i,0], label=f'Particle {i+1}') + ax3.plot(tracing.times, 299792458*trajectory[:, 3]/jnp.sqrt(tracing.energy()[i]/mass*2.), label=f'Particle {i+1}') ax4.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') - - -ax2.set_xlabel('Time (s)') +ax2.set_xlabel(r'$t~[\mathrm{s}]$') ax2.set_ylabel('Normalized energy variation') ax3.set_ylabel(r'$v_{\parallel}/v$') -#ax2.legend() -ax3.set_xlabel('Time (s)') -#ax3.legend() -ax4.set_xlabel('R (m)') -ax4.set_ylabel('Z (m)') -#ax4.legend() +ax3.set_xlabel(r'$t~[\mathrm{s}]$') +ax4.set_xlabel(r'$R~[\mathrm{m}]$') +ax4.set_ylabel(r'$Z~[\mathrm{m}]$') plt.tight_layout() -plt.savefig('traj.pdf') +plt.savefig('traj_time.pdf') - -v=jnp.sqrt(tracing.energy*2./particles.mass) -#pitch=trajectories[:,:,3]/v -vpar=trajectories[:,:,3] -vperp=tracing.vperp_final +v=jnp.sqrt(tracing.energy()*2./particles.mass) +vpar=trajectories[:,:,3]*SPEED_OF_LIGHT +vpar=jnp.where(jnp.isfinite(vpar), vpar, jnp.nan) +vperp=tracing.v_perp() pitch=vpar/v # Plot distribution in velocities initial t and final fig3 = plt.figure(figsize=(9, 8)) @@ -198,7 +177,6 @@ plt.savefig('statistics.pdf') - # Plot distribution in velocities initial t and final fig2 = plt.figure(figsize=(9, 8)) ax12 = fig2.add_subplot(251) @@ -212,12 +190,12 @@ ax92 = fig2.add_subplot(259) nbins=64 -v0=jnp.sqrt(tracing.energy[:,0]*2./particles.mass) -vfinal=jnp.sqrt(tracing.energy[:,-1]*2./particles.mass) -vperp0=tracing.vperp_final[:,0] -vperpfinal=tracing.vperp_final[:,-1] -vpar0=trajectories[:,0,3] -vparfinal=trajectories[:,-1,3] +v0=jnp.sqrt(tracing.energy()[:,0]*2./particles.mass)/SPEED_OF_LIGHT +vfinal=jnp.sqrt(tracing.energy()[:,-1]*2./particles.mass)/SPEED_OF_LIGHT +vperp0=tracing.v_perp()[:,0]/SPEED_OF_LIGHT +vperpfinal=tracing.v_perp()[:,-1]/SPEED_OF_LIGHT +vpar0=vpar[:,0]/SPEED_OF_LIGHT +vparfinal=vpar[:,-1]/SPEED_OF_LIGHT pitch0=vpar0/v0 pitch_final=vparfinal/vfinal @@ -339,7 +317,16 @@ def find_first_less_than_numpy(arr, value): ax92.set_xlabel(r'$t_{final}$') plt.tight_layout() -plt.savefig('dist.pdf') + +# Improved time distribution plot +plt.figure(figsize=(7, 5)) +plt.hist(good_t_final, bins=nbins, color='c', edgecolor='black', alpha=0.7) +plt.title(r'$t_{final}$ Distribution', fontweight='bold') +plt.xlabel(r'$t_{final}$', fontweight='bold') +plt.ylabel('Counts', fontweight='bold') +plt.tight_layout() +plt.rcParams.update({'font.size': 18, 'font.weight': 'bold'}) +plt.savefig('dist_time.pdf', dpi=300) ## Save results in vtk format to analyze in Paraview # tracing.to_vtk('trajectories') diff --git a/examples/trace_particles_coils_guidingcenter_with_classifier_with_collisionsMu.py b/examples/particle_tracing_collisions/trace_particles_coils_guidingcenter_with_classifier_with_collisionsMu.py similarity index 84% rename from examples/trace_particles_coils_guidingcenter_with_classifier_with_collisionsMu.py rename to examples/particle_tracing_collisions/trace_particles_coils_guidingcenter_with_classifier_with_collisionsMu.py index 1e6cbf71..95a01cf4 100644 --- a/examples/trace_particles_coils_guidingcenter_with_classifier_with_collisionsMu.py +++ b/examples/particle_tracing_collisions/trace_particles_coils_guidingcenter_with_classifier_with_collisionsMu.py @@ -7,12 +7,11 @@ import matplotlib.pyplot as plt from essos.fields import BiotSavart,Vmec from essos.surfaces import SurfaceClassifier -from essos.coils import Coils_from_json,Coils_from_simsopt +from essos.coils import Coils from essos.constants import ALPHA_PARTICLE_MASS, ALPHA_PARTICLE_CHARGE, FUSION_ALPHA_PARTICLE_ENERGY,ONE_EV,ELECTRON_MASS,PROTON_MASS,SPEED_OF_LIGHT from essos.dynamics import Tracing, Particles from essos.background_species import BackgroundSpecies - # Input parameters tmax = 1e-4 timestep=1.e-8 @@ -24,7 +23,6 @@ rtol=1.e-5 energy=FUSION_ALPHA_PARTICLE_ENERGY - #Initialize background species #number_species=2 #(electrons,deuterium) #mass_array=jnp.array([ELECTRON_MASS/PROTON_MASS,2]) #mass_over_mproton @@ -42,18 +40,18 @@ species = BackgroundSpecies(number_species=number_species, mass_array=mass_array, charge_array=charge_array, n_array=n_array, T_array=T_array) # Load coils and field -#json_file = os.path.join(os.path.dirname(__name__), 'input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') -#coils = Coils_from_json(json_file) +#json_file = os.path.join(os.path.dirname(__file__), '..', 'input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') +#coils = Coils.from_json(json_file) #field = BiotSavart(coils) # Load coils and field -json_file = os.path.join(os.path.dirname(__name__), 'input_files', 'QH_simple_scaled.json')#'SIMSOPT_biot_savart_LandremanPaulQA.json') -coils = Coils_from_simsopt(json_file,nfp=4) +json_file = os.path.join(os.path.dirname(__file__), '..', 'input_files', 'QH_simple_scaled.json')#'SIMSOPT_biot_savart_LandremanPaulQA.json') +coils = Coils.from_simsopt(json_file,nfp=4) field = BiotSavart(coils) # Load coils and field -wout_file = os.path.join(os.path.dirname(__name__), 'input_files','wout_QH_simple_scaled.nc') +wout_file = os.path.join(os.path.dirname(__file__), '..', 'input_files','wout_QH_simple_scaled.nc') vmec = Vmec(wout_file) timeI=time() @@ -89,8 +87,8 @@ for i, trajectory in enumerate(trajectories): #ax2.plot(tracing.times, jnp.abs(tracing.energy[i]-particles.energy)/particles.energy, label=f'Particle {i+1}') - ax2.plot(tracing.times, (tracing.energy[i]-tracing.energy[i][0])/particles.energy, label=f'Particle {i+1}') - ax3.plot(tracing.times, trajectory[:, 3]*SPEED_OF_LIGHT/particles.total_speed, label=f'Particle {i+1}') + ax2.plot(tracing.times, (tracing.energy()[i]-tracing.energy()[i][0])/tracing.energy()[i][0], label=f'Particle {i+1}') + ax3.plot(tracing.times, trajectory[:, 3]*SPEED_OF_LIGHT/tracing.energy()[i], label=f'Particle {i+1}') #ax4.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') ax4.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') ax2.set_xlabel('Time (s)') diff --git a/examples/trace_particles_vmec_collisionsMu.py b/examples/particle_tracing_collisions/trace_particles_vmec_collisionsMu.py similarity index 91% rename from examples/trace_particles_vmec_collisionsMu.py rename to examples/particle_tracing_collisions/trace_particles_vmec_collisionsMu.py index 6d00464b..a5edb6d6 100644 --- a/examples/trace_particles_vmec_collisionsMu.py +++ b/examples/particle_tracing_collisions/trace_particles_vmec_collisionsMu.py @@ -10,7 +10,6 @@ from essos.background_species import BackgroundSpecies import numpy as np - # Input parameters tmax = 1.e-4 timestep=1.e-8 @@ -27,7 +26,7 @@ energy=FUSION_ALPHA_PARTICLE_ENERGY # Load coils and field -wout_file = os.path.join(os.path.dirname(__name__), 'input_files',"wout_LandremanPaul2021_QA_reactorScale_lowres.nc") +wout_file = os.path.join(os.path.dirname(__file__), '..', 'input_files',"wout_LandremanPaul2021_QA_reactorScale_lowres.nc") vmec = Vmec(wout_file) # Initialize particles @@ -71,9 +70,9 @@ for i in np.random.choice(nparticles, size=n_particles_to_plot, replace=False): trajectory = trajectories[i] ## Plot energy error - ax2.plot(tracing.times, (tracing.energy[i]-tracing.energy[i][0])/tracing.energy[i][0], label=f'Particle {i+1}') + ax2.plot(tracing.times, (tracing.energy()[i]-tracing.energy()[i,0])/tracing.energy()[i,0], label=f'Particle {i+1}') ## Plot velocity parallel to the magnetic field - ax3.plot(tracing.times, trajectory[:, 3]*SPEED_OF_LIGHT/jnp.sqrt(tracing.energy[i]/particles.mass*2.), label=f'Particle {i+1}') + ax3.plot(tracing.times, trajectory[:, 3]*SPEED_OF_LIGHT/jnp.sqrt(tracing.energy()[i]/particles.mass*2.), label=f'Particle {i+1}') ## Plot s-coordinate ax4.plot(tracing.times, trajectory[:,0], label=f'Particle {i+1}') # ax4.set_ylabel(r'$s=\psi/\psi_b$') diff --git a/examples/coils_from_BOOZ_XFORM.py b/examples/simple_examples/coils_from_BOOZ_XFORM.py similarity index 100% rename from examples/coils_from_BOOZ_XFORM.py rename to examples/simple_examples/coils_from_BOOZ_XFORM.py diff --git a/examples/coils_from_nearaxis.py b/examples/simple_examples/coils_from_nearaxis.py similarity index 98% rename from examples/coils_from_nearaxis.py rename to examples/simple_examples/coils_from_nearaxis.py index 8518d231..9e67a12f 100644 --- a/examples/coils_from_nearaxis.py +++ b/examples/simple_examples/coils_from_nearaxis.py @@ -4,7 +4,9 @@ import jax.numpy as jnp from jax import block_until_ready, vmap import matplotlib.pyplot as plt -from essos.fields import near_axis, BiotSavart_from_gamma, BiotSavart +from essos.fields import BiotSavart_from_gamma, BiotSavart +from pyqsc_jax.near_axis import near_axis + import plotly.graph_objects as go from essos.dynamics import Tracing from essos.coils import fit_dofs_from_coils, Curves, Coils diff --git a/examples/create_perturbed_coils.py b/examples/simple_examples/create_perturbed_coils.py similarity index 79% rename from examples/create_perturbed_coils.py rename to examples/simple_examples/create_perturbed_coils.py index b5109ad5..f0e8ee17 100644 --- a/examples/create_perturbed_coils.py +++ b/examples/simple_examples/create_perturbed_coils.py @@ -10,11 +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 # Coils parameters order_Fourier_series_coils = 4 @@ -35,21 +31,21 @@ -g=GaussianSampler(coils_initial.quadpoints,sigma=0.2,length_scale=0.1,n_derivs=2) +g=GaussianSampler(coils_initial.curves.quadpoints,sigma=0.2,length_scale=0.1,n_derivs=2) #Split the key for reproducibility key=0 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, 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, 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, g, key=split_keys[0]) -perturb_curves_statistic(coils_perturbed, 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)) @@ -61,13 +57,11 @@ plt.legend() plt.show() - - # # Save the coils to a json file # coils_optimized.to_json("stellarator_coils.json") # # Load the coils from a json file -# from essos.coils import Coils_from_json -# coils = Coils_from_json("stellarator_coils.json") +# from essos.coils import Coils +# coils = Coils.from_json("stellarator_coils.json") # # Save results in vtk format to analyze in Paraview # tracing_initial.to_vtk('trajectories_initial') diff --git a/examples/create_stellarator_coils.py b/examples/simple_examples/create_stellarator_coils.py similarity index 92% rename from examples/create_stellarator_coils.py rename to examples/simple_examples/create_stellarator_coils.py index 304e4f3b..812fa8be 100644 --- a/examples/create_stellarator_coils.py +++ b/examples/simple_examples/create_stellarator_coils.py @@ -32,8 +32,8 @@ # # Save the coils to a json file # coils.to_json("stellarator_coils.json") # # Load the coils from a json file -# from essos.coils import Coils_from_json -# coils = Coils_from_json("stellarator_coils.json") +# from essos.coils import Coils +# coils = Coils.from_json("stellarator_coils.json") # # View coils in Paraview # coils.to_vtk('stellarator_coils') \ No newline at end of file diff --git a/examples/simple_examples/get_derivatives_coils_particle_confinement_guidingcenter.py b/examples/simple_examples/get_derivatives_coils_particle_confinement_guidingcenter.py new file mode 100644 index 00000000..baedacb7 --- /dev/null +++ b/examples/simple_examples/get_derivatives_coils_particle_confinement_guidingcenter.py @@ -0,0 +1,56 @@ + +import os +from jax import vmap +import jax.numpy as jnp +from essos.coils import Coils, CreateEquallySpacedCurves +from essos.fields import BiotSavart +from essos.losses import custom_loss + +# Optimization parameters +order_Fourier_series_coils = 2 +number_coil_points = 80 +number_coils_per_half_field_period = 3 +number_of_field_periods = 2 + +LENGTH_TARGET = 31 +CURVATURE_TARGET = 0.4 +AXIS_B_TARGET = 5.7 + +# Initialize coils +current_on_each_coil = 1.84e7 +major_radius_coils = 7.75 +minor_radius_coils = 4.45 +curves = CreateEquallySpacedCurves(n_curves=number_coils_per_half_field_period, + order=order_Fourier_series_coils, + R=major_radius_coils, r=minor_radius_coils, + n_segments=number_coil_points, + nfp=number_of_field_periods, stellsym=True) +coils_initial = Coils(curves=curves, currents=[current_on_each_coil]*number_coils_per_half_field_period) +field = BiotSavart(coils_initial) + +""" Creating the loss functions """ +def loss_length(field): + return jnp.mean(jnp.maximum(0, field.coils.length - LENGTH_TARGET)) + +def loss_curvature(field): + return jnp.mean(jnp.maximum(0, field.coils.curvature - CURVATURE_TARGET)) + +def loss_normB_axis_average(field): + R_axis=field.r_axis + phi_array = jnp.linspace(0, 2 * jnp.pi, 15) + B_axis = vmap(lambda phi: field.AbsB(jnp.array([R_axis * jnp.cos(phi), R_axis * jnp.sin(phi), 0])))(phi_array) + return jnp.mean(jnp.maximum(0, B_axis - AXIS_B_TARGET)) + +curvature_loss = custom_loss(loss_curvature, "field") +length_loss = custom_loss(loss_length, "field") +Baxis_average_loss = custom_loss(loss_normB_axis_average, "field") +total_loss = curvature_loss + length_loss + Baxis_average_loss +total_loss.dependencies = {"field": field} + +## Take the gradients +params = total_loss.starting_dofs +loss = total_loss(params) +gradients = total_loss.grad(params) + +print('Objective function: {:.2E}'.format(loss)) +print('Gradients (derivative of objective function with respect to coils): ',gradients) diff --git a/examples/testing_collisions_velocity_distributions_mu_Adaptative.py b/examples/testing_collisions_velocity_distributions_mu_Adaptative.py deleted file mode 100644 index 3a19ed84..00000000 --- a/examples/testing_collisions_velocity_distributions_mu_Adaptative.py +++ /dev/null @@ -1,274 +0,0 @@ -import os -number_of_processors_to_use = 1 # Parallelization, this should divide nparticles -os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' -from time import time -import jax.numpy as jnp -import matplotlib.pyplot as plt -import matplotlib.colors -from essos.fields import BiotSavart -from essos.coils import Coils_from_json -from essos.constants import PROTON_MASS, ONE_EV,ELECTRON_MASS,SPEED_OF_LIGHT -from essos.dynamics import Tracing, Particles -from essos.background_species import BackgroundSpecies,gamma_ab -import numpy as np -import jax - -from jax import config -# to use higher precision -config.update("jax_enable_x64", True) - - - - -# Input parameters -tmax = 1e-5 -dt=1.e-14 -times_to_trace=100 -nparticles_per_core=10 -nparticles = number_of_processors_to_use*nparticles_per_core -R0 = 1.25#jnp.linspace(1.23, 1.27, nparticles) -atol = 1.e-6 -rtol=0. -rejected_steps=100 -mass=PROTON_MASS -mass_e=ELECTRON_MASS -T_test=3000. -energy=T_test*ONE_EV - - -# Load coils and field -json_file = os.path.join(os.path.dirname(__name__), 'input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') -coils = Coils_from_json(json_file) -field = BiotSavart(coils) - -# Initialize particles -Z0 = jnp.zeros(nparticles) -phi0 = jnp.zeros(nparticles) -initial_xyz=jnp.array([R0*jnp.cos(phi0), R0*jnp.sin(phi0), Z0]).T -particles = Particles(initial_xyz=initial_xyz,initial_vparallel_over_v=1.0*jnp.ones(nparticles), mass=mass, energy=energy) - - - - -#Initialize background species -number_species=1 #(electrons,deuterium) -mass_array=jnp.array([1.]) #mass_over_mproton -charge_array=jnp.array([1.]) #mass_over_mproton -T0=1.e+3 #eV -n0=1e+20 #m^-3 -n_array=jnp.array([n0]) -T_array=jnp.array([T0]) -species = BackgroundSpecies(number_species=number_species, mass_array=mass_array, charge_array=charge_array, n_array=n_array, T_array=T_array) -vth_c=jnp.sqrt(T0*ONE_EV/PROTON_MASS)/SPEED_OF_LIGHT -vpar_mean=0. -vpar_sigma=vth_c -v_mean=vth_c*jnp.sqrt(8./jnp.pi) -v_sigma=vth_c*jnp.sqrt((3.*jnp.pi-8.)/jnp.pi) -vperp_mean=vth_c*jnp.sqrt(jnp.pi/2.) -vperp_sigma=vth_c*jnp.sqrt(2.-jnp.pi/2.) -pitch_mean=0. -pitch_sigma=jnp.sqrt(2.**2/12) - - -# Trace in ESSOS -time0 = time() -tracing = Tracing(field=field, model='GuidingCenterCollisionsMuAdaptative', particles=particles, - maxtime=tmax, timestep=dt,times_to_trace=times_to_trace, rtol=rtol,atol=atol,species=species,tag_gc=0.,rejected_steps=100) -print(f"ESSOS tracing took {time()-time0:.2f} seconds") -trajectories = tracing.trajectories - -# Plot trajectories, velocity parallel to the magnetic field, and energy error -fig = plt.figure(figsize=(9, 8)) -ax1 = fig.add_subplot(221, projection='3d') -ax2 = fig.add_subplot(222) -ax3 = fig.add_subplot(223) -ax4 = fig.add_subplot(224) - -coils.plot(ax=ax1, show=False) -tracing.plot(ax=ax1, show=False) - - -for i, trajectory in enumerate(trajectories): - ax2.plot(tracing.times, (tracing.energy[i]-tracing.energy[i,0])/tracing.energy[i,0], label=f'Particle {i+1}') - ax3.plot(tracing.times, trajectory[:, 3]*SPEED_OF_LIGHT/jnp.sqrt(tracing.energy[i]/mass*2.), label=f'Particle {i+1}') - ax4.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') - - - - -ax2.set_xlabel('Time (s)') -ax2.set_ylabel('Normalized energy variation') -ax3.set_ylabel(r'$v_{\parallel}/v$') -ax3.set_xlabel('Time (s)') -ax4.set_xlabel('R (m)') -ax4.set_ylabel('Z (m)') -plt.tight_layout() -plt.savefig('traj.pdf') - - - - -v=jnp.sqrt(tracing.energy*2./particles.mass) -vpar=trajectories[:,:,3]*SPEED_OF_LIGHT -vperp=tracing.vperp_final -pitch=vpar/v -# Plot distribution in velocities initial t and final -fig3 = plt.figure(figsize=(9, 8)) -ax13 = fig3.add_subplot(241) -ax23 = fig3.add_subplot(242) -ax33 = fig3.add_subplot(243) -ax43 = fig3.add_subplot(244) -ax53 = fig3.add_subplot(245) -ax63 = fig3.add_subplot(246) -ax73 = fig3.add_subplot(247) -ax83 = fig3.add_subplot(248) -ax13.plot(tracing.times,jnp.nanmean(v/SPEED_OF_LIGHT,axis=0)) -ax13.axhline(y=v_mean, color='r', linestyle='--') -ax23.plot(tracing.times,jnp.nanstd(v/SPEED_OF_LIGHT,axis=0)) -ax23.axhline(y=v_sigma, color='r', linestyle='--') -ax33.plot(tracing.times,jnp.nanmean(pitch,axis=0)) -ax33.axhline(y=pitch_mean, color='r', linestyle='--') -ax43.plot(tracing.times,jnp.nanstd(pitch,axis=0)) -ax43.axhline(y=pitch_sigma, color='r', linestyle='--') -ax53.plot(tracing.times,jnp.nanmean(vpar/SPEED_OF_LIGHT,axis=0)) -ax53.axhline(y=vpar_mean, color='r', linestyle='--') -ax63.plot(tracing.times,jnp.nanstd(vpar/SPEED_OF_LIGHT,axis=0)) -ax63.axhline(y=vpar_sigma, color='r', linestyle='--') -ax73.plot(tracing.times,jnp.nanmean(vperp/SPEED_OF_LIGHT,axis=0)) -ax73.axhline(y=vperp_mean, color='r', linestyle='--') -ax83.plot(tracing.times,jnp.nanstd(vperp/SPEED_OF_LIGHT,axis=0)) -ax83.axhline(y=vperp_sigma, color='r', linestyle='--') -ax13.set_title('Mean energy') -ax13.set_xlabel('time') -ax13.set_ylabel('Energy') -ax23.set_title('sigma energy') -ax23.set_xlabel('time') -ax23.set_ylabel('Energy') -ax33.set_title('Mean pitch') -ax33.set_xlabel('time') -ax33.set_ylabel('pitch') -ax43.set_title('sigma pitch') -ax43.set_xlabel('time') -ax43.set_ylabel('pitch') -ax53.set_title('Mean vpar') -ax53.set_xlabel('time') -ax53.set_ylabel('vpar') -ax63.set_title('sigma vpar') -ax63.set_xlabel('time') -ax63.set_ylabel('vpar') -ax73.set_title('Mean vperp') -ax73.set_xlabel('time') -ax73.set_ylabel('vperp') -ax83.set_title('sigma vperp') -ax83.set_xlabel('time') -ax83.set_ylabel('vperp') -plt.tight_layout() -plt.savefig('statistics.pdf') - - - -# Plot distribution in velocities initial t and final -fig2 = plt.figure(figsize=(9, 8)) -ax12 = fig2.add_subplot(251) -ax22 = fig2.add_subplot(252) -ax32 = fig2.add_subplot(253) -ax42 = fig2.add_subplot(254) -ax52 = fig2.add_subplot(255) -ax62 = fig2.add_subplot(256) -ax72 = fig2.add_subplot(257) -ax82 = fig2.add_subplot(258) -nbins=64 - -v0=jnp.sqrt(tracing.energy[:,0]*2./particles.mass) -vfinal=jnp.sqrt(tracing.energy[:,-1]*2./particles.mass) -vperp0=tracing.vperp_final[:,0] -vperpfinal=tracing.vperp_final[:,-1] -vpar0=trajectories[:,0,3] -vparfinal=trajectories[:,-1,3] -pitch0=vpar0/v0 -pitch_final=vparfinal/vfinal - - - - - -bad_indices_v0 = jnp.isnan(v0) -bad_indices_vfinal = jnp.isnan(vfinal) -bad_indices_pitch0 = jnp.isnan(pitch0) -bad_indices_pitch_final = jnp.isnan(pitch_final) -bad_indices_vperp0 = jnp.isnan(vperp0) -bad_indices_vperp_final = jnp.isnan(vperpfinal) -bad_indices_vpar0 = jnp.isnan(vpar0) -bad_indices_vpar_final = jnp.isnan(vparfinal) -good_indices_v0 = ~bad_indices_v0 -good_indices_vfinal = ~bad_indices_vfinal -good_indices_pitch0 = ~bad_indices_pitch0 -good_indices_pitch_final = ~bad_indices_pitch_final -good_indices_vpar0 = ~bad_indices_vpar0 -good_indices_vpar_final = ~bad_indices_vpar_final -good_indices_vperp0 = ~bad_indices_vperp0 -good_indices_vperp_final = ~bad_indices_vperp_final -good_v0 = v0[good_indices_v0] -good_vfinal = vfinal[good_indices_vfinal] -good_pitch0 = pitch0[good_indices_pitch0] -good_pitch_final = pitch_final[good_indices_pitch_final] - - -good_vpar0 = vpar0[good_indices_vpar0] -good_vpar_final = vparfinal[good_indices_vpar_final] -good_vperp0 = vperp0[good_indices_vperp0] -good_vperp_final = vperpfinal[good_indices_vperp_final] - - -v0_counts,v0_bins=jnp.histogram(good_v0,bins=nbins) -vfinal_counts,vfinal_bins=jnp.histogram(good_vfinal,bins=nbins) - -pitch_t0_counts,pitch_t0_bins=jnp.histogram(good_pitch0,bins=nbins) -pitch_tfinal_counts,pitch_tfinal_bins=jnp.histogram(good_pitch_final,bins=nbins) - -vpar_t0_counts,vpar_t0_bins=jnp.histogram(good_vpar0,bins=nbins) -vpar_tfinal_counts,vpar_tfinal_bins=jnp.histogram(good_vpar_final,bins=nbins) - - -vperp_t0_counts,vperp_t0_bins=jnp.histogram(good_vperp0,bins=nbins) -vperp_tfinal_counts,vperp_tfinal_bins=jnp.histogram(good_vperp_final,bins=nbins) - - -ax12.stairs(v0_counts,v0_bins) -ax22.stairs(vfinal_counts,vfinal_bins) -ax32.stairs(vpar_t0_counts,vpar_t0_bins) -ax42.stairs(vpar_tfinal_counts,vpar_tfinal_bins) -ax52.stairs(pitch_t0_counts,pitch_t0_bins) -ax62.stairs(pitch_tfinal_counts,pitch_tfinal_bins) -ax72.stairs(vperp_t0_counts,vperp_t0_bins) -ax82.stairs(vperp_tfinal_counts,vperp_tfinal_bins) - -ax12.set_title('t=0') -ax12.set_xlabel('v') -ax12.set_ylabel('Counts') -ax22.set_title('t=t_final') -ax22.set_xlabel('v') -ax22.set_ylabel('Counts') -ax32.set_title('t=0') -ax32.set_ylabel('Counts') -ax32.set_xlabel(r'$v_{\parallel}$') -ax42.set_title('t=t_final') -ax42.set_xlabel(r'$v_{parallel}$') -ax42.set_ylabel('Counts') -ax52.set_title('t=0') -ax52.set_xlabel(r'$v_{\parallel}/v$') -ax52.set_ylabel('Counts') -ax62.set_title('t=t_final') -ax62.set_xlabel(r'$v_{\parallel}/v$') -ax62.set_ylabel('Counts') -ax72.set_title('t=0') -ax72.set_ylabel('Counts') -ax72.set_xlabel(r'$v_{\perp}$') -ax82.set_title('t=t_final') -ax82.set_ylabel('Counts') -ax82.set_xlabel(r'$v_{\perp}$') - - -plt.tight_layout() -plt.savefig('dist.pdf') - diff --git a/examples/testing_collisions_velocity_distributions_mu_Fixed.py b/examples/testing_collisions_velocity_distributions_mu_Fixed.py deleted file mode 100644 index 4c268807..00000000 --- a/examples/testing_collisions_velocity_distributions_mu_Fixed.py +++ /dev/null @@ -1,266 +0,0 @@ -import os -number_of_processors_to_use = 1 # Parallelization, this should divide nparticles -os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}' -from time import time -import jax.numpy as jnp -import matplotlib.pyplot as plt -import matplotlib.colors -from essos.fields import BiotSavart -from essos.coils import Coils_from_json -from essos.constants import PROTON_MASS, ONE_EV,ELECTRON_MASS,SPEED_OF_LIGHT -from essos.dynamics import Tracing, Particles -from essos.background_species import BackgroundSpecies,gamma_ab -import numpy as np -import jax - -from jax import config -# to use higher precision -config.update("jax_enable_x64", True) - -# Input parameters -tmax = 1.e-5 -dt=1.e-8 -times_to_trace=100 -nparticles_per_core=10 -nparticles = number_of_processors_to_use*nparticles_per_core -R0 = 1.25 -num_steps = jnp.round(tmax/dt) -mass=PROTON_MASS -mass_e=ELECTRON_MASS -T_test=3000. -energy=T_test*ONE_EV - - -# Load coils and field -json_file = os.path.join(os.path.dirname(__name__), 'input_files', 'ESSOS_biot_savart_LandremanPaulQA.json') -coils = Coils_from_json(json_file) -field = BiotSavart(coils) - -# Initialize particles -Z0 = jnp.zeros(nparticles) -phi0 = jnp.zeros(nparticles) -initial_xyz=jnp.array([R0*jnp.cos(phi0), R0*jnp.sin(phi0), Z0]).T -particles = Particles(initial_xyz=initial_xyz,initial_vparallel_over_v=1.0*jnp.ones(nparticles), mass=mass, energy=energy) - - -#Initialize background species -number_species=1 #(electrons,deuterium) -mass_array=jnp.array([1.]) #mass_over_mproton -charge_array=jnp.array([1.]) #mass_over_mproton -T0=1.e+3 #eV -n0=1e+20 #m^-3 -n_array=jnp.array([n0]) -T_array=jnp.array([T0]) -species = BackgroundSpecies(number_species=number_species, mass_array=mass_array, charge_array=charge_array, n_array=n_array, T_array=T_array) -vth_c=jnp.sqrt(T0*ONE_EV/PROTON_MASS)/SPEED_OF_LIGHT -vpar_mean=0. -vpar_sigma=vth_c -v_mean=vth_c*jnp.sqrt(8./jnp.pi) -v_sigma=vth_c*jnp.sqrt((3.*jnp.pi-8.)/jnp.pi) -vperp_mean=vth_c*jnp.sqrt(jnp.pi/2.) -vperp_sigma=vth_c*jnp.sqrt(2.-jnp.pi/2.) -pitch_mean=0. -pitch_sigma=jnp.sqrt(2.**2/12) - - -# Trace in ESSOS -time0 = time() -tracing = Tracing(field=field, model='GuidingCenterCollisionsMuFixed', particles=particles, - maxtime=tmax, timestep=dt,times_to_trace=times_to_trace,species=species,tag_gc=0.) -print(f"ESSOS tracing took {time()-time0:.2f} seconds") -trajectories = tracing.trajectories - -# Plot trajectories, velocity parallel to the magnetic field, and energy error -fig = plt.figure(figsize=(9, 8)) -ax1 = fig.add_subplot(221, projection='3d') -ax2 = fig.add_subplot(222) -ax3 = fig.add_subplot(223) -ax4 = fig.add_subplot(224) - -coils.plot(ax=ax1, show=False) -tracing.plot(ax=ax1, show=False) - -for i, trajectory in enumerate(trajectories): - ax2.plot(tracing.times, (tracing.energy[i]-tracing.energy[i,0])/tracing.energy[i,0], label=f'Particle {i+1}') - ax3.plot(tracing.times, 299792458*trajectory[:, 3]/jnp.sqrt(tracing.energy[i]/mass*2.), label=f'Particle {i+1}') - ax4.plot(jnp.sqrt(trajectory[:,0]**2+trajectory[:,1]**2), trajectory[:, 2], label=f'Particle {i+1}') - - - - -ax2.set_xlabel('Time (s)') -ax2.set_ylabel('Normalized energy variation') -ax3.set_ylabel(r'$v_{\parallel}/v$') -ax3.set_xlabel('Time (s)') -ax4.set_xlabel('R (m)') -ax4.set_ylabel('Z (m)') -plt.tight_layout() -plt.savefig('traj.pdf') - - - - -v=jnp.sqrt(tracing.energy*2./particles.mass) -vpar=trajectories[:,:,3] -vperp=tracing.vperp_final -pitch=vpar/v - -# Plot distribution in velocities initial t and final -fig3 = plt.figure(figsize=(9, 8)) -ax13 = fig3.add_subplot(241) -ax23 = fig3.add_subplot(242) -ax33 = fig3.add_subplot(243) -ax43 = fig3.add_subplot(244) -ax53 = fig3.add_subplot(245) -ax63 = fig3.add_subplot(246) -ax73 = fig3.add_subplot(247) -ax83 = fig3.add_subplot(248) -ax13.plot(tracing.times,jnp.nanmean(v/SPEED_OF_LIGHT,axis=0)) -ax13.axhline(y=v_mean, color='r', linestyle='--') -ax23.plot(tracing.times,jnp.nanstd(v/SPEED_OF_LIGHT,axis=0)) -ax23.axhline(y=v_sigma, color='r', linestyle='--') -ax33.plot(tracing.times,jnp.nanmean(pitch,axis=0)) -ax33.axhline(y=pitch_mean, color='r', linestyle='--') -ax43.plot(tracing.times,jnp.nanstd(pitch,axis=0)) -ax43.axhline(y=pitch_sigma, color='r', linestyle='--') -ax53.plot(tracing.times,jnp.nanmean(vpar/SPEED_OF_LIGHT,axis=0)) -ax53.axhline(y=vpar_mean, color='r', linestyle='--') -ax63.plot(tracing.times,jnp.nanstd(vpar/SPEED_OF_LIGHT,axis=0)) -ax63.axhline(y=vpar_sigma, color='r', linestyle='--') -ax73.plot(tracing.times,jnp.nanmean(vperp/SPEED_OF_LIGHT,axis=0)) -ax73.axhline(y=vperp_mean, color='r', linestyle='--') -ax83.plot(tracing.times,jnp.nanstd(vperp/SPEED_OF_LIGHT,axis=0)) -ax83.axhline(y=vperp_sigma, color='r', linestyle='--') -ax13.set_title('Mean energy') -ax13.set_xlabel('time') -ax13.set_ylabel('Energy') -ax23.set_title('sigma energy') -ax23.set_xlabel('time') -ax23.set_ylabel('Energy') -ax33.set_title('Mean pitch') -ax33.set_xlabel('time') -ax33.set_ylabel('pitch') -ax43.set_title('sigma pitch') -ax43.set_xlabel('time') -ax43.set_ylabel('pitch') -ax53.set_title('Mean vpar') -ax53.set_xlabel('time') -ax53.set_ylabel('vpar') -ax63.set_title('sigma vpar') -ax63.set_xlabel('time') -ax63.set_ylabel('vpar') -ax73.set_title('Mean vperp') -ax73.set_xlabel('time') -ax73.set_ylabel('vperp') -ax83.set_title('sigma vperp') -ax83.set_xlabel('time') -ax83.set_ylabel('vperp') -plt.tight_layout() -plt.savefig('statistics.pdf') - - - -# Plot distribution in velocities initial t and final -fig2 = plt.figure(figsize=(9, 8)) -ax12 = fig2.add_subplot(251) -ax22 = fig2.add_subplot(252) -ax32 = fig2.add_subplot(253) -ax42 = fig2.add_subplot(254) -ax52 = fig2.add_subplot(255) -ax62 = fig2.add_subplot(256) -ax72 = fig2.add_subplot(257) -ax82 = fig2.add_subplot(258) -nbins=64 - -v0=jnp.sqrt(tracing.energy[:,0]*2./particles.mass) -vfinal=jnp.sqrt(tracing.energy[:,-1]*2./particles.mass) -vperp0=tracing.vperp_final[:,0] -vperpfinal=tracing.vperp_final[:,-1] -vpar0=trajectories[:,0,3] -vparfinal=trajectories[:,-1,3] -pitch0=vpar0/v0 -pitch_final=vparfinal/vfinal - - - - - -bad_indices_v0 = jnp.isnan(v0) -bad_indices_vfinal = jnp.isnan(vfinal) -bad_indices_pitch0 = jnp.isnan(pitch0) -bad_indices_pitch_final = jnp.isnan(pitch_final) -bad_indices_vperp0 = jnp.isnan(vperp0) -bad_indices_vperp_final = jnp.isnan(vperpfinal) -bad_indices_vpar0 = jnp.isnan(vpar0) -bad_indices_vpar_final = jnp.isnan(vparfinal) -good_indices_v0 = ~bad_indices_v0 -good_indices_vfinal = ~bad_indices_vfinal -good_indices_pitch0 = ~bad_indices_pitch0 -good_indices_pitch_final = ~bad_indices_pitch_final -good_indices_vpar0 = ~bad_indices_vpar0 -good_indices_vpar_final = ~bad_indices_vpar_final -good_indices_vperp0 = ~bad_indices_vperp0 -good_indices_vperp_final = ~bad_indices_vperp_final -good_v0 = v0[good_indices_v0] -good_vfinal = vfinal[good_indices_vfinal] -good_pitch0 = pitch0[good_indices_pitch0] -good_pitch_final = pitch_final[good_indices_pitch_final] - - -good_vpar0 = vpar0[good_indices_vpar0] -good_vpar_final = vparfinal[good_indices_vpar_final] -good_vperp0 = vperp0[good_indices_vperp0] -good_vperp_final = vperpfinal[good_indices_vperp_final] - - -v0_counts,v0_bins=jnp.histogram(good_v0,bins=nbins) -vfinal_counts,vfinal_bins=jnp.histogram(good_vfinal,bins=nbins) - -pitch_t0_counts,pitch_t0_bins=jnp.histogram(good_pitch0,bins=nbins) -pitch_tfinal_counts,pitch_tfinal_bins=jnp.histogram(good_pitch_final,bins=nbins) - -vpar_t0_counts,vpar_t0_bins=jnp.histogram(good_vpar0,bins=nbins) -vpar_tfinal_counts,vpar_tfinal_bins=jnp.histogram(good_vpar_final,bins=nbins) - - -vperp_t0_counts,vperp_t0_bins=jnp.histogram(good_vperp0,bins=nbins) -vperp_tfinal_counts,vperp_tfinal_bins=jnp.histogram(good_vperp_final,bins=nbins) - - -ax12.stairs(v0_counts,v0_bins) -ax22.stairs(vfinal_counts,vfinal_bins) -ax32.stairs(vpar_t0_counts,vpar_t0_bins) -ax42.stairs(vpar_tfinal_counts,vpar_tfinal_bins) -ax52.stairs(pitch_t0_counts,pitch_t0_bins) -ax62.stairs(pitch_tfinal_counts,pitch_tfinal_bins) -ax72.stairs(vperp_t0_counts,vperp_t0_bins) -ax82.stairs(vperp_tfinal_counts,vperp_tfinal_bins) - -ax12.set_title('t=0') -ax12.set_xlabel('v') -ax12.set_ylabel('Counts') -ax22.set_title('t=t_final') -ax22.set_xlabel('v') -ax22.set_ylabel('Counts') -ax32.set_title('t=0') -ax32.set_ylabel('Counts') -ax32.set_xlabel(r'$v_{\parallel}$') -ax42.set_title('t=t_final') -ax42.set_xlabel(r'$v_{parallel}$') -ax42.set_ylabel('Counts') -ax52.set_title('t=0') -ax52.set_xlabel(r'$v_{\parallel}/v$') -ax52.set_ylabel('Counts') -ax62.set_title('t=t_final') -ax62.set_xlabel(r'$v_{\parallel}/v$') -ax62.set_ylabel('Counts') -ax72.set_title('t=0') -ax72.set_ylabel('Counts') -ax72.set_xlabel(r'$v_{\perp}$') -ax82.set_title('t=t_final') -ax82.set_ylabel('Counts') -ax82.set_xlabel(r'$v_{\perp}$') - - -plt.tight_layout() -plt.savefig('dist.pdf') diff --git a/pyproject.toml b/pyproject.toml index 1f770eb1..4ce0f4ec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,8 +27,20 @@ keywords = ["Plasma", "Simulation", "JAX"] -dependencies = [ "jax", "jaxlib", "tqdm", "matplotlib", "diffrax", "optax", "jaxopt", "optimistix", "scipy", "jaxkd", "netcdf4"] - +dependencies = [ + "jax", + "jaxlib", + "tqdm", + "matplotlib", + "diffrax", + "optax", + "jaxopt", + "optimistix", + "scipy", + "jaxkd", + "netcdf4", + "pyqsc_jax @ git+https://github.com/uwplasma/pyQSC_JAX.git" +] requires-python = ">=3.10" [project.urls] diff --git a/requirements.txt b/requirements.txt index aea84879..989b7a13 100644 --- a/requirements.txt +++ b/requirements.txt @@ -12,4 +12,4 @@ netcdf4 f90nml pyevtk optuna -pandas \ No newline at end of file +pandas diff --git a/tests/test_augmented_lagrangian.py b/tests/test_augmented_lagrangian.py index 6be5c41d..b91658de 100644 --- a/tests/test_augmented_lagrangian.py +++ b/tests/test_augmented_lagrangian.py @@ -11,6 +11,7 @@ eq, ineq, combine, + SelectiveConstraint, total_infeasibility, norm_constraints, infty_norm_constraints, @@ -28,14 +29,16 @@ class TestAugmentedLagrangian(unittest.TestCase): def test_lagrange_multiplier(self): - lm = LagrangeMultiplier(value=1.0, penalty=2.0, sq_grad=3.0) + lm = LagrangeMultiplier(value=1.0, penalty=2.0, omega=4.0, eta=5.0, sq_grad=3.0) self.assertEqual(lm.value, 1.0) self.assertEqual(lm.penalty, 2.0) + self.assertEqual(lm.omega, 4.0) + self.assertEqual(lm.eta, 5.0) self.assertEqual(lm.sq_grad, 3.0) def test_update_method_all_modes(self): - params = LagrangeMultiplier(jnp.array([1.]), jnp.array([2.]), jnp.array([0.])) - updates = LagrangeMultiplier(jnp.array([0.5]), jnp.array([0.]), jnp.array([0.])) + params = LagrangeMultiplier(value=jnp.array([1.]), penalty=jnp.array([2.]), omega=jnp.array([0.]), eta=jnp.array([0.]), sq_grad=jnp.array([0.])) + updates = LagrangeMultiplier(value=jnp.array([0.5]), penalty=jnp.array([0.]), omega=jnp.array([0.]), eta=jnp.array([0.]), sq_grad=jnp.array([0.])) for mode in [ 'Constant', 'Mu_Monotonic', 'Mu_Conditional_True', 'Mu_Conditional_False', 'Mu_Tolerance_True', 'Mu_Tolerance_False', 'Mu_Adaptative' @@ -50,8 +53,8 @@ def test_update_method_all_modes(self): self.assertIsInstance(result, LagrangeMultiplier) def test_update_method_squared_all_modes(self): - params = LagrangeMultiplier(jnp.array([1.]), jnp.array([2.]), jnp.array([0.])) - updates = LagrangeMultiplier(jnp.array([0.5]), jnp.array([0.]), jnp.array([0.])) + params = LagrangeMultiplier(value=jnp.array([1.]), penalty=jnp.array([2.]), omega=jnp.array([0.]), eta=jnp.array([0.]), sq_grad=jnp.array([0.])) + updates = LagrangeMultiplier(value=jnp.array([0.5]), penalty=jnp.array([0.]), omega=jnp.array([0.]), eta=jnp.array([0.]), sq_grad=jnp.array([0.])) for mode in [ 'Constant', 'Mu_Monotonic', 'Mu_Conditional_True', 'Mu_Conditional_False', 'Mu_Tolerance_True', 'Mu_Tolerance_False', 'Mu_Adaptative' @@ -103,6 +106,51 @@ def fun3(x): return x * 2 params = combined.init(jnp.array([2.])) combined.loss(params, jnp.array([2.])) + def test_selective_constraint_dependencies_reset_cached_dofs(self): + selective = SelectiveConstraint(eq(lambda field: field - 1), 'field') + selective.dependencies = {'field': jnp.array([1.0, 2.0])} + first = selective.starting_dofs + + selective.dependencies = {'field': jnp.array([3.0, 4.0, 5.0])} + second = selective.starting_dofs + + self.assertEqual(first.shape[0], 2) + self.assertEqual(second.shape[0], 3) + self.assertTrue(jnp.allclose(second, jnp.array([3.0, 4.0, 5.0]))) + + def test_composite_constraint_dependencies_reset_cached_dofs(self): + c1 = SelectiveConstraint(eq(lambda field: field - 1), 'field') + c2 = SelectiveConstraint(eq(lambda surface: surface + 1), 'surface') + combined = combine(c1, c2) + combined.dependencies = { + 'field': jnp.array([1.0, 2.0]), + 'surface': jnp.array([10.0]), + } + first = combined.starting_dofs + + combined.dependencies = { + 'field': jnp.array([3.0]), + 'surface': jnp.array([20.0, 30.0]), + } + second = combined.starting_dofs + + self.assertEqual(first.shape[0], 3) + self.assertEqual(second.shape[0], 3) + self.assertTrue(jnp.allclose(second, jnp.array([3.0, 20.0, 30.0]))) + + def test_composite_constraint_set_dependencies_resets_cached_dofs(self): + c1 = SelectiveConstraint(eq(lambda field: field - 1), 'field') + combined = combine(c1) + combined.set_dependencies({'field': jnp.array([1.0, 2.0])}) + first = combined.starting_dofs + + combined.set_dependencies({'field': jnp.array([7.0])}) + second = combined.starting_dofs + + self.assertEqual(first.shape[0], 2) + self.assertEqual(second.shape[0], 1) + self.assertTrue(jnp.allclose(second, jnp.array([7.0]))) + def test_total_infeasibility(self): tree = {'a': jnp.array([1.0, -2.0]), 'b': jnp.array([3.0])} result = total_infeasibility(tree) @@ -119,7 +167,7 @@ def test_infty_norm_constraints(self): self.assertAlmostEqual(float(result), 3.0) def test_penalty_average(self): - tree = {'a': LagrangeMultiplier(jnp.array([1.0]), jnp.array([2.0]), jnp.array([0.0]))} + tree = {'a': LagrangeMultiplier(value=jnp.array([1.0]), penalty=jnp.array([2.0]), omega=jnp.array([0.0]), eta=jnp.array([0.0]), sq_grad=jnp.array([0.0]))} result = penalty_average(tree) self.assertAlmostEqual(float(result), 2.0) @@ -144,8 +192,8 @@ def test_lagrange_update_gradient_transformation_and_update(self): self.assertTrue(hasattr(gt, 'update')) # Call init and update with dummy data params = {'x': jnp.array([1.0])} - lagrange_params = LagrangeMultiplier(jnp.array([0.0]), jnp.array([1.0]), jnp.array([0.0])) - updates = LagrangeMultiplier(jnp.array([-0.5]), jnp.array([1.0]), jnp.array([1.0])) + lagrange_params = LagrangeMultiplier(value=jnp.array([0.0]), penalty=jnp.array([1.0]), omega=jnp.array([0.0]), eta=jnp.array([0.0]), sq_grad=jnp.array([0.0])) + updates = LagrangeMultiplier(value=jnp.array([-0.5]), penalty=jnp.array([1.0]), omega=jnp.array([1.0]), eta=jnp.array([1.0]), sq_grad=jnp.array([1.0])) state = gt.init(params) # eta, omega, etc. are required by update_fn signature eta = {'x': jnp.array([0.0])} @@ -244,4 +292,4 @@ def fun(x): return x - 1 if __name__ == "__main__": - pytest.main([__file__]) \ No newline at end of file + pytest.main([__file__]) diff --git a/tests/test_coil_perturbation.py b/tests/test_coil_perturbation.py index a33821e6..4405ddbe 100644 --- a/tests/test_coil_perturbation.py +++ b/tests/test_coil_perturbation.py @@ -108,7 +108,7 @@ def test_perturb_curves_systematic(self): key = jax.random.PRNGKey(0) for sampler in [sampler0, sampler1, sampler2]: curves = DummyCurves(n_base_curves=2, nfp=1, stellsym=True, n_points=5) - perturb_curves_systematic(curves, sampler, key) + curves = perturb_curves_systematic(curves, sampler, key) # Just check that gamma arrays are still the right shape self.assertEqual(curves.gamma.shape, (2, 5, 3)) @@ -120,8 +120,8 @@ def test_perturb_curves_statistic(self): key = jax.random.PRNGKey(0) for sampler in [sampler0, sampler1, sampler2]: curves = DummyCurves(n_base_curves=2, nfp=1, stellsym=True, n_points=5) - perturb_curves_statistic(curves, sampler, key) + curves = perturb_curves_statistic(curves, sampler, key) self.assertEqual(curves.gamma.shape, (2, 5, 3)) if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main() 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_multiobjectives.py b/tests/test_multiobjectives.py index 019a0ac5..d8333762 100644 --- a/tests/test_multiobjectives.py +++ b/tests/test_multiobjectives.py @@ -1,9 +1,10 @@ -import pytest -from unittest.mock import MagicMock, patch -from essos.multiobjectiveoptimizer import MultiObjectiveOptimizer -from essos.coils import Coils,Curves -from essos.fields import BiotSavart -from essos.objective_functions import loss_bdotn_over_b, loss_coil_length, loss_coil_curvature, loss_normB_axis +import pytest +from unittest.mock import MagicMock, patch +from essos.losses import custom_loss +from essos.multiobjectiveoptimizer import MultiObjectiveOptimizer +from essos.coils import Coils,Curves +from essos.fields import BiotSavart +from essos.objective_functions import loss_bdotn_over_b, loss_coil_length, loss_coil_curvature, loss_normB_axis # test_multiobjectiveoptimizer.py @@ -28,15 +29,48 @@ def mock_vmec(): -def dummy_loss_fn(): - def loss_fn(field=None, coils=None, vmec=None, surface=None, x=None): - return jnp.sum(x) - return loss_fn - - -def test_build_available_inputs( vmec=mock_vmec(), dummy_loss_fn=dummy_loss_fn()): - optimizer = MultiObjectiveOptimizer( - loss_functions=[dummy_loss_fn], +def dummy_loss_fn(field=None, coils=None, vmec=None, surface=None, x=None): + return jnp.sum(x) + + +def test_custom_loss_named_unraveler(): + def loss_fn(curve_dofs, current): + return jnp.sum(curve_dofs**2) + jnp.sum(current) + + loss = custom_loss(loss_fn, "curve_dofs", "current") + loss.dependencies = { + "curve_dofs": jnp.array([[1.0, 2.0], [3.0, 4.0]]), + "current": jnp.array([5.0]), + "unused": jnp.array([99.0]), + } + + dofs = loss.starting_dofs + named_args = loss.dofs_to_pytree(dofs) + tuple_args = tuple(named_args[name] for name in loss.args_names) + + assert set(named_args) == {"curve_dofs", "current"} + assert jnp.array_equal(named_args["curve_dofs"], loss.dependencies["curve_dofs"]) + assert jnp.array_equal(named_args["current"], loss.dependencies["current"]) + assert loss(dofs) == loss_fn(named_args["curve_dofs"], named_args["current"]) + assert loss.call_pytree(named_args) == loss_fn(named_args["curve_dofs"], named_args["current"]) + assert loss.call_pytree(tuple_args) == loss_fn(named_args["curve_dofs"], named_args["current"]) + value, grad = loss.value_and_grad(dofs) + assert value == loss_fn(named_args["curve_dofs"], named_args["current"]) and jnp.array_equal(grad, loss.grad(dofs)) + + gradient = loss.grad_pytree(named_args) + gradient_tuple = loss.grad_pytree(tuple_args) + assert set(gradient) == {"curve_dofs", "current", "unused"} + assert jnp.array_equal(gradient["curve_dofs"], 2 * named_args["curve_dofs"]) + assert jnp.array_equal(gradient["current"], jnp.ones_like(named_args["current"])) + assert jnp.array_equal(gradient["unused"], jnp.zeros_like(loss.dependencies["unused"])) + assert jnp.array_equal(gradient_tuple["curve_dofs"], gradient["curve_dofs"]) + assert jnp.array_equal(gradient_tuple["current"], gradient["current"]) + assert jnp.array_equal(gradient_tuple["unused"], gradient["unused"]) + + +def test_build_available_inputs( vmec=mock_vmec(), dummy_loss_fn=dummy_loss_fn): + optimizer = MultiObjectiveOptimizer( + loss_functions=[dummy_loss_fn], vmec=vmec, coils_init=None, function_inputs={"extra": 42}, diff --git a/tests/test_objective_functions.py b/tests/test_objective_functions.py index c2d6eedc..0bec8389 100644 --- a/tests/test_objective_functions.py +++ b/tests/test_objective_functions.py @@ -1,261 +1,275 @@ import unittest +from types import SimpleNamespace from unittest.mock import MagicMock, patch + +import jax import jax.numpy as jnp import essos.objective_functions as objf + class DummyField: def __init__(self): - self.R0 = jnp.array([1.]) - self.Z0 = jnp.array([0.]) - self.phi = jnp.array([0.]) - self.B_axis = jnp.array([[1., 0., 0.]]) - self.grad_B_axis = jnp.array([[0., 0., 0.]]) + self.R0 = jnp.array([1.0, 1.0]) + self.Z0 = jnp.array([0.0, 0.0]) + self.phi = jnp.array([0.0, jnp.pi / 2]) + self.B_axis = jnp.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + self.grad_B_axis = jnp.ones((2, 3, 3)) + self.iota = 0.4 + self.elongation = jnp.array([1.0, 1.2]) self.r_axis = 1.0 self.z_axis = 0.0 - self.AbsB = MagicMock(return_value=5.7) - self.B = MagicMock(return_value=jnp.array([1., 0., 0.])) - self.dB_by_dX = MagicMock(return_value=jnp.array([0., 0., 0.])) - self.B_covariant = MagicMock(return_value=jnp.array([1., 0., 0.])) - self.coils_length = jnp.array([30.]) - self.coils_curvature = jnp.ones((2, 10)) - self.gamma = jnp.zeros((2, 10, 3)) - self.gamma_dash = jnp.ones((2, 10, 3)) - self.gamma_dashdash = jnp.ones((2, 10, 3)) - self.currents = jnp.ones(2) - self.quadpoints = jnp.linspace(0, 1, 10) - self.x = jnp.zeros((10,1)) - -class DummyCoils(DummyField): - def __init__(self): - super().__init__() + self.coils = DummyCoils() + + def AbsB(self, points): + return 5.7 + 0.1 * jnp.sum(points) + + def B(self, points): + return jnp.array([1.0 + 0.1 * points[0], 0.5, 0.25]) + + def dB_by_dX(self, points): + return jnp.eye(3) + + def B_covariant(self, points): + return jnp.array([1.0, 0.5, 0.25]) + + def copy(self): + return DummyField() -class DummyCurves: - def __init__(self, *args, **kwargs): - pass class DummyParticles: def __init__(self): - self.to_full_orbit = MagicMock() - self.trajectories = jnp.zeros((2, 10, 3)) + self.energy = 1.0 + self.mass = 1.0 + self.charge = 1.0 + self.total_speed = 1.0 + + def to_full_orbit(self, field): + return None + class DummyTracing: def __init__(self, *args, **kwargs): - self.trajectories = jnp.zeros((2, 10, 3)) - self.field = DummyField() - self.loss_fractions = jnp.array([0.1,0.2,1.]) - self.times_to_trace = 10 + xyz = jnp.array( + [ + [[1.0, 0.0, 0.0], [1.1, 0.0, 0.05], [1.2, 0.0, 0.1], [1.3, 0.0, 0.15]], + [[1.0, 0.0, 0.0], [1.05, 0.0, 0.04], [1.1, 0.0, 0.08], [1.15, 0.0, 0.12]], + ], + dtype=jnp.float64, + ) + self.trajectories = xyz + self.field = kwargs.get("field", DummyField()) + self.loss_fractions = jnp.array([0.1, 0.2, 1.0]) + self.times_to_trace = 4 self.maxtime = 1e-5 -class DummyVmec: - def __init__(self): - self.surface = MagicMock() class DummySurface: def __init__(self): - self.gamma = jnp.zeros((10, 3)) - self.unitnormal = jnp.ones((10, 3)) - -def dummy_sampler(*args, **kwargs): - return 0 - -def dummy_new_nearaxis_from_x_and_old_nearaxis(x, field_nearaxis): - class DummyNearAxis: - elongation = jnp.array([1.]) - iota = 1.0 - x = jnp.array([1.]) - R0 = jnp.array([1.]) - Z0 = jnp.array([0.]) - phi = jnp.array([0.]) - B_axis = jnp.array([[1., 0., 0.]]) - grad_B_axis = jnp.array([[0., 0., 0.]]) - return DummyNearAxis() + self.gamma = jnp.zeros((2, 3, 3), dtype=jnp.float64) + self.unitnormal = jnp.ones((2, 3, 3), dtype=jnp.float64) + self.stellsym = False + self.nfp = 1 + + +@jax.tree_util.register_pytree_node_class +class PytreeCoils: + def __init__(self, gamma, gamma_dash, gamma_dashdash, currents, quadpoints, length, curvature, base_curves, order=1, nfp=1, stellsym=False): + self.gamma = gamma + self.gamma_dash = gamma_dash + self.gamma_dashdash = gamma_dashdash + self.currents = currents + self.length = length + self.curvature = curvature + self.order = order + self.nfp = nfp + self.stellsym = stellsym + self.curves = SimpleNamespace(quadpoints=quadpoints, curves=base_curves) + + def __len__(self): + return self.gamma.shape[0] + + def copy(self): + return PytreeCoils( + self.gamma, + self.gamma_dash, + self.gamma_dashdash, + self.currents, + self.curves.quadpoints, + self.length, + self.curvature, + self.curves.curves, + order=self.order, + nfp=self.nfp, + stellsym=self.stellsym, + ) + + def tree_flatten(self): + children = ( + self.gamma, + self.gamma_dash, + self.gamma_dashdash, + self.currents, + self.curves.quadpoints, + self.length, + self.curvature, + self.curves.curves, + ) + aux = {"order": self.order, "nfp": self.nfp, "stellsym": self.stellsym} + return children, aux + + @classmethod + def tree_unflatten(cls, aux, children): + return cls(*children, **aux) + + +@jax.tree_util.register_pytree_node_class +class PytreeSurface: + def __init__(self, gamma, unitnormal, stellsym=False, nfp=1): + self.gamma = gamma + self.unitnormal = unitnormal + self.stellsym = stellsym + self.nfp = nfp + + def tree_flatten(self): + return (self.gamma, self.unitnormal), (self.stellsym, self.nfp) + + @classmethod + def tree_unflatten(cls, aux, children): + gamma, unitnormal = children + stellsym, nfp = aux + return cls(gamma, unitnormal, stellsym=stellsym, nfp=nfp) + class TestObjectiveFunctions(unittest.TestCase): def setUp(self): - self.x = jnp.ones(12) - self.dofs_curves = jnp.ones((2, 3)) - self.currents_scale = 1.0 - self.nfp = 1 - self.n_segments = 10 - self.stellsym = True - self.key = 0 - self.sampler = dummy_sampler self.field = DummyField() - self.coils = DummyCoils() - self.curves = DummyCurves() + self.field_nearaxis = DummyField() self.particles = DummyParticles() - self.tracing = DummyTracing() - self.vmec = DummyVmec() self.surface = DummySurface() - - @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): - 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) - - @patch('essos.objective_functions.Curves', return_value=DummyCurves()) - @patch('essos.objective_functions.Coils', return_value=DummyCoils()) - @patch('essos.objective_functions.BiotSavart', return_value=DummyField()) - def test_field_and_coils_from_dofs(self, bs, coils, curves): - objf.field_from_dofs(self.x, self.dofs_curves, self.currents_scale, self.nfp) - objf.coils_from_dofs(self.x, self.dofs_curves, self.currents_scale, self.nfp) - objf.curves_from_dofs(self.x, self.dofs_curves, self.nfp) - - @patch('essos.objective_functions.field_from_dofs', return_value=DummyField()) - def test_loss_coil_length_and_curvature(self, ffd): - objf.loss_coil_length(self.x, self.dofs_curves, self.currents_scale, self.nfp) - objf.loss_coil_curvature(self.x, self.dofs_curves, self.currents_scale, self.nfp) - objf.loss_coil_length_new(self.x, self.dofs_curves, self.currents_scale, self.nfp) - objf.loss_coil_curvature_new(self.x, self.dofs_curves, self.currents_scale, self.nfp) - - @patch('essos.objective_functions.field_from_dofs', return_value=DummyField()) - def test_loss_normB_axis(self, ffd): - objf.loss_normB_axis(self.x, self.dofs_curves, self.currents_scale, self.nfp) - objf.loss_normB_axis_average(self.x, self.dofs_curves, self.currents_scale, self.nfp) - - @patch('essos.objective_functions.field_from_dofs', return_value=DummyField()) - def test_loss_particle_functions(self, ffd): - with patch('essos.objective_functions.Tracing', return_value=self.tracing): - objf.loss_particle_radial_drift(self.x, self.particles, self.dofs_curves, self.currents_scale, self.nfp) - objf.loss_particle_alpha_drift(self.x, self.particles, self.dofs_curves, self.currents_scale, self.nfp) - objf.loss_particle_gamma_c(self.x, self.particles, self.dofs_curves, self.currents_scale, self.nfp) - objf.loss_particle_r_cross_final(self.x, self.particles, self.dofs_curves, self.currents_scale, self.nfp) - objf.loss_particle_r_cross_max_constraint(self.x, self.particles, self.dofs_curves, self.currents_scale, self.nfp) - objf.loss_Br(self.x, self.particles, self.dofs_curves, self.currents_scale, self.nfp) - objf.loss_iota(self.x, self.particles, self.dofs_curves, self.currents_scale, self.nfp) - - @patch('essos.objective_functions.field_from_dofs', return_value=DummyField()) - def test_loss_lost_fraction(self, ffd): - with patch('essos.objective_functions.Tracing', return_value=self.tracing): - objf.loss_lost_fraction(self.field, self.particles, self.dofs_curves, self.currents_scale, self.nfp) - - def test_normB_axis(self): - objf.normB_axis(self.field) - - @patch('essos.objective_functions.field_from_dofs', return_value=DummyField()) - @patch('essos.objective_functions.new_nearaxis_from_x_and_old_nearaxis', side_effect=dummy_new_nearaxis_from_x_and_old_nearaxis) - def test_loss_coils_for_nearaxis_and_loss_coils_and_nearaxis(self, nna, ffd): - objf.loss_coils_for_nearaxis(self.x, self.field, self.dofs_curves, self.currents_scale, self.nfp) - objf.loss_coils_and_nearaxis(jnp.ones(13), self.field, self.dofs_curves, self.currents_scale, self.nfp) - - def test_difference_B_gradB_onaxis(self): - objf.difference_B_gradB_onaxis(self.field, self.field) - - @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.BdotN_over_B', return_value=jnp.ones(10)) - def test_loss_bdotn_over_b(self, bdotn, bs, coils, curves): - objf.loss_bdotn_over_b(self.x, self.vmec, self.dofs_curves, self.currents_scale, self.nfp) - - @patch('essos.objective_functions.field_from_dofs', return_value=DummyField()) - @patch('essos.objective_functions.BdotN_over_B', return_value=jnp.ones(10)) - def test_loss_BdotN(self, bdotn, ffd): - objf.loss_BdotN(self.x, self.vmec, self.dofs_curves, self.currents_scale, self.nfp) - - @patch('essos.objective_functions.field_from_dofs', return_value=DummyField()) - @patch('essos.objective_functions.BdotN_over_B', return_value=jnp.ones(10)) - def test_loss_BdotN_only(self, bdotn, ffd): - objf.loss_BdotN_only(self.x, self.vmec, self.dofs_curves, self.currents_scale, self.nfp) - - @patch('essos.objective_functions.field_from_dofs', return_value=DummyField()) - @patch('essos.objective_functions.BdotN_over_B', return_value=jnp.ones(10)) - def test_loss_BdotN_only_constraint(self, bdotn, ffd): - objf.loss_BdotN_only_constraint(self.x, self.vmec, self.dofs_curves, self.currents_scale, self.nfp) - - @patch('essos.objective_functions.BdotN_over_B', return_value=jnp.ones(10)) - @patch('essos.objective_functions.pertubred_field_from_dofs', return_value=DummyField()) - def test_loss_BdotN_only_stochastic(self, perturbed, bdotn): - objf.loss_BdotN_only_stochastic(self.x, self.sampler, 2, self.vmec, self.dofs_curves, self.currents_scale, self.nfp) - - @patch('essos.objective_functions.BdotN_over_B', return_value=jnp.ones(10)) - @patch('essos.objective_functions.pertubred_field_from_dofs', return_value=DummyField()) - def test_loss_BdotN_only_constraint_stochastic(self, perturbed, bdotn): - objf.loss_BdotN_only_constraint_stochastic(self.x, self.sampler, 2, self.vmec, self.dofs_curves, self.currents_scale, self.nfp) - - @patch('essos.objective_functions.coils_from_dofs', return_value=DummyCoils()) - def test_loss_cs_distance_and_array(self, cfd): - objf.loss_cs_distance(self.x, self.surface, self.dofs_curves, self.currents_scale, self.nfp) - objf.loss_cs_distance_array(self.x, self.surface, self.dofs_curves, self.currents_scale, self.nfp) - - @patch('essos.objective_functions.coils_from_dofs', return_value=DummyCoils()) - def test_loss_cc_distance_and_array(self, cfd): - objf.loss_cc_distance(self.x, self.dofs_curves, self.currents_scale, self.nfp) - objf.loss_cc_distance_array(self.x, self.dofs_curves, self.currents_scale, self.nfp) - - @patch('essos.objective_functions.coils_from_dofs', return_value=DummyCoils()) - def test_loss_linking_mnumber_and_constraint(self, cfd): - objf.loss_linking_mnumber(self.x, self.dofs_curves, self.currents_scale, self.nfp) - objf.loss_linking_mnumber_constarint(self.x, self.dofs_curves, self.currents_scale, self.nfp) - - def test_cc_distance_pure(self): - gamma1 = jnp.ones((10, 3))*3. - l1 = jnp.ones((10, 3)) - gamma2 = jnp.ones((10, 3))*4. - l2 = jnp.ones((10, 3))*6. - objf.cc_distance_pure(gamma1, l1, gamma2, l2, 1.0) - - def test_cs_distance_pure(self): - gammac = jnp.ones((10, 3))*7. - lc = jnp.ones((10, 3)) - gammas = jnp.ones((10, 3))*9. - ns = jnp.ones((10, 3))*10. - objf.cs_distance_pure(gammac, lc, gammas, ns, 1.0) - - @patch('essos.objective_functions.coils_from_dofs', return_value=DummyCoils()) - def test_loss_lorentz_force_coils(self, cfd): - objf.loss_lorentz_force_coils(self.x, self.dofs_curves, self.currents_scale, self.nfp) - - @patch('essos.objective_functions.compute_curvature', return_value=1.0) - @patch('essos.objective_functions.BiotSavart_from_gamma', return_value=MagicMock(B=MagicMock(return_value=jnp.array([1., 0., 0.])))) - def test_lp_force_pure(self, bsg, cc): - gamma = jnp.ones((2, 10, 3))*2. - gamma_dash = jnp.ones((2, 10, 3))*3. - gamma_dashdash = jnp.ones((2, 10, 3)) - currents = jnp.ones(2) - quadpoints = jnp.linspace(0, 1, 10) - objf.lp_force_pure(0, gamma, gamma_dash, gamma_dashdash, currents, quadpoints, 1, 1e6) - - def test_B_regularized_singularity_term(self): + self.sampler = MagicMock(name="sampler") + self.keys = jnp.array([0, 1], dtype=jnp.int32) + self.coils = PytreeCoils( + gamma=jnp.arange(2 * 5 * 3, dtype=jnp.float64).reshape(2, 5, 3) / 10.0, + gamma_dash=jnp.ones((2, 5, 3), dtype=jnp.float64), + gamma_dashdash=jnp.ones((2, 5, 3), dtype=jnp.float64) * 0.1, + currents=jnp.ones(2, dtype=jnp.float64), + quadpoints=jnp.linspace(0.0, 1.0, 5), + length=jnp.array([3.0, 4.0], dtype=jnp.float64), + curvature=jnp.ones((2, 5), dtype=jnp.float64), + base_curves=jnp.arange(2 * 3 * 3, dtype=jnp.float64).reshape(2, 3, 3) / 10.0, + ) + self.pytree_surface = PytreeSurface( + gamma=jnp.arange(2 * 3 * 3, dtype=jnp.float64).reshape(2, 3, 3) / 20.0, + unitnormal=jnp.ones((2, 3, 3), dtype=jnp.float64), + ) + + def test_near_axis_losses(self): + points, B_nearaxis, gradB_nearaxis = objf.near_axis_field_quantities(self.field_nearaxis) + self.assertEqual(points.shape, (3, 2)) + self.assertEqual(B_nearaxis.shape, (2, 3)) + self.assertEqual(gradB_nearaxis.shape, (3, 3, 2)) + self.assertTrue(jnp.isfinite(objf.loss_B_difference_coils_near_axis(self.field, self.field_nearaxis))) + self.assertTrue(jnp.isfinite(objf.loss_gradB_difference_coils_near_axis(self.field, self.field_nearaxis))) + self.assertTrue(jnp.isfinite(objf.loss_iota_near_axis(self.field_nearaxis))) + self.assertTrue(jnp.isfinite(objf.loss_r0_near_axis(self.field_nearaxis))) + + @patch("essos.objective_functions.Tracing", side_effect=DummyTracing) + def test_particle_losses(self, tracing): + self.assertTrue(jnp.isfinite(objf.loss_particle_radial_drift(self.field, self.particles))) + self.assertTrue(jnp.isfinite(objf.loss_particle_radial_drift_fullorbit(self.field, self.particles))) + self.assertTrue(jnp.isfinite(objf.loss_particle_alpha_drift(self.field, self.particles))) + self.assertTrue(jnp.isfinite(objf.loss_particle_gammac(self.field, self.particles))) + self.assertTrue(jnp.isfinite(objf.loss_particle_rcross_final(self.field, self.particles))) + self.assertTrue(jnp.isfinite(objf.loss_particle_Br(self.field, self.particles))) + self.assertTrue(jnp.isfinite(objf.loss_particle_iota(self.field, self.particles))) + + @patch("essos.objective_functions.BdotN_over_B", return_value=jnp.ones((2, 3), dtype=jnp.float64)) + def test_surface_losses(self, bdotn): + self.assertTrue(jnp.isfinite(objf.normB_axis(self.field)).all()) + self.assertTrue(jnp.isfinite(objf.loss_normB_axis_average(self.field))) + self.assertTrue(jnp.isfinite(objf.loss_BdotN(self.field, self.surface))) + self.assertTrue(jnp.isfinite(objf.loss_BdotN_constraint(self.field, self.surface))) + + def test_copy_coils_from_field(self): + copied = objf.copy_coils_from_field(self.field) + self.assertIsInstance(copied, 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_from_field(self, statistical, systematic, biot_savart): + systematic.side_effect = lambda coils, sampler, key=None: coils + statistical.side_effect = lambda coils, sampler, key=None: coils + field = objf.perturbed_field_from_field(self.field, 0, self.sampler) + self.assertIsInstance(field, DummyField) + systematic.assert_called_once() + statistical.assert_called_once() + + @patch("essos.objective_functions.BdotN_over_B", return_value=jnp.ones((2, 3), dtype=jnp.float64)) + @patch("essos.objective_functions.perturbed_field_from_field", return_value=DummyField()) + def test_stochastic_surface_losses(self, perturbed_field, bdotn): + self.assertTrue(jnp.isfinite(objf.loss_bdotn_stochastic(self.field, self.surface, self.sampler, self.keys))) + self.assertTrue(jnp.isfinite(objf.constraint_bdotn_stochastic(self.field, self.surface, self.sampler, self.keys))) + + def test_coil_length_and_curvature_losses(self): + self.assertTrue(jnp.isfinite(objf.loss_coil_length(self.coils, max_coil_length=4.0)).all()) + self.assertTrue(jnp.isfinite(objf.loss_coil_curvature(self.coils, max_coil_curvature=2.0)).all()) + + def test_compute_candidates(self): + i_vals, j_vals = objf.compute_candidates(self.coils, min_separation=10.0) + self.assertEqual(i_vals.ndim, 1) + self.assertEqual(j_vals.ndim, 1) + + def test_blockwise_losses_with_non_divisible_blocks(self): + separation = objf.loss_coil_separation(self.coils, 0.5, block_size=3) + surface_distance = objf.loss_coil_surface_distance(self.coils, self.pytree_surface, 0.5, block_size=4) + linking = objf.loss_linkingnumber(self.coils, block_size=4) + self.assertTrue(jnp.isfinite(separation)) + self.assertTrue(jnp.isfinite(surface_distance)) + self.assertTrue(jnp.isfinite(linking)) + self.assertAlmostEqual(float(separation), float(objf.loss_coil_separation.__wrapped__(self.coils, 0.5, block_size=3))) + self.assertAlmostEqual(float(surface_distance), float(objf.loss_coil_surface_distance.__wrapped__(self.coils, self.pytree_surface, 0.5, block_size=4))) + self.assertAlmostEqual(float(linking), float(objf.loss_linkingnumber.__wrapped__(self.coils, block_size=4))) + + @patch("essos.objective_functions.Curves.compute_curvature", return_value=jnp.ones(5)) + @patch("essos.objective_functions.BiotSavart_from_gamma") + def test_loss_lorentz_force_coils(self, biot_savart_from_gamma, compute_curvature): + class DummyBS: + def B(self, point): + return jnp.zeros(3) + + biot_savart_from_gamma.return_value = DummyBS() + force_loss = objf.loss_lorentz_force_coils(self.coils, threshold=1e6, block_size=2) + self.assertTrue(jnp.isfinite(force_loss)) + self.assertAlmostEqual( + float(force_loss), + float(objf.loss_lorentz_force_coils.__wrapped__(self.coils, threshold=1e6, block_size=2)), + ) + + def test_regularization_helpers(self): rc_prime = jnp.ones((10, 3)) rc_prime_prime = jnp.ones((10, 3)) - objf.B_regularized_singularity_term(rc_prime, rc_prime_prime, 1.0) - - def test_B_regularized_pure(self): - gamma = jnp.ones((10, 3))*4. + gamma = jnp.ones((10, 3)) * 4.0 gammadash = jnp.ones((10, 3)) gammadashdash = jnp.ones((10, 3)) quadpoints = jnp.linspace(0, 1, 10) - current = 1.0 - regularization = 1.0 - objf.B_regularized_pure(gamma, gammadash, gammadashdash, quadpoints, current, regularization) - - def test_regularization_circ(self): + self.assertTrue(jnp.isfinite(objf.B_regularized_singularity_term(rc_prime, rc_prime_prime, 1.0)).all()) + self.assertTrue(jnp.isfinite(objf.B_regularized_pure(gamma, gammadash, gammadashdash, quadpoints, 1.0, 1.0)).all()) self.assertTrue(objf.regularization_circ(2.0) > 0) + self.assertTrue(jnp.isfinite(objf.regularization_rect(2.0, 1.0))) + self.assertTrue(jnp.isfinite(objf.rectangular_xsection_k(2.0, 1.0))) + self.assertTrue(jnp.isfinite(objf.rectangular_xsection_delta(2.0, 1.0))) + + +class DummyCoils: + def __init__(self): + self.length = jnp.array([3.0, 4.0]) + self.curvature = jnp.ones((2, 5)) + + def copy(self): + return DummyCoils() - def test_regularization_rect_and_k_and_delta(self): - a, b = 2.0, 1.0 - objf.regularization_rect(a, b) - objf.rectangular_xsection_k(a, b) - objf.rectangular_xsection_delta(a, b) - - def test_linking_number_pure_and_integrand(self): - gamma1 = jnp.ones((10, 3))*4. - lc1 = jnp.ones((10, 3))*2. - gamma2 = jnp.ones((10, 3))*6. - lc2 = jnp.ones((10, 3))*5. - dphi = 0.1 - objf.linking_number_pure(gamma1, lc1, gamma2, lc2, dphi) - r1 = jnp.zeros(3) - dr1 = jnp.ones(3) - r2 = jnp.zeros(3) - dr2 = jnp.ones(3) - objf.integrand_linking_number(r1, dr1, r2, dr2, dphi, dphi) if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main()