Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@ to compute the toroidal flux when possible, as opposed to a 2D surface integral
- if an ``int``, this is the chunk size to be used.
- if ``"auto"`` for the ``ObjectiveFunction``, will use a heuristic for the maximum ``jac_chunk_size`` needed to fit the jacobian calculation on the available device memory, according to the formula: ``max_jac_chunk_size = (desc_config.get("avail_mem") / estimated_memory_usage - 0.22) / 0.85 * self.dim_x`` with ``estimated_memory_usage = 2.4e-7 * self.dim_f * self.dim_x + 1``
- the ``ObjectiveFunction`` ``jac_chunk_size`` is used if ``deriv_mode="batched"``, and the ``_Objective`` ``jac_chunk_size`` will be used if ``deriv_mode="blocked"``
- Add ``from_input_file`` method to ``Equilibrium`` class to generate an ``Equilibrium`` object with boundary, profiles, resolution and flux specified in a given DESC or VMEC input file



Bug Fixes
Expand Down
51 changes: 51 additions & 0 deletions desc/equilibrium/equilibrium.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import copy
import numbers
import warnings
from collections.abc import MutableSequence

import numpy as np
Expand All @@ -26,6 +27,7 @@
ZernikeRZToroidalSection,
)
from desc.grid import Grid, LinearGrid, QuadratureGrid, _Grid
from desc.input_reader import InputReader
from desc.io import IOAble
from desc.objectives import (
ForceBalance,
Expand Down Expand Up @@ -2035,6 +2037,55 @@ def from_near_axis(

return eq

@classmethod
def from_input_file(cls, path, **kwargs):
"""Create an Equilibrium from information in a DESC or VMEC input file.

Parameters
----------
path : Path-like or str
Path to DESC or VMEC input file.
**kwargs : dict, optional
keyword arguments to pass to the constructor of the
Equilibrium being created.

Returns
-------
Equilibrium : Equilibrium
Equilibrium generated from the given input file.

"""
inputs = InputReader().parse_inputs(path)[-1]
if (inputs["bdry_ratio"] is not None) and (inputs["bdry_ratio"] != 1):
warnings.warn(
"`bdry_ratio` is intended as an input for the continuation method."
"`bdry_ratio`=1 uses the given surface modes as is, any other scalar "
"value will scale the non-axisymmetric modes by that value. The "
"final value of `bdry_ratio` in the input file is "
f"{inputs['bdry_ratio']}, this means the created Equilibrium won't "
"have the given surface but a scaled version instead."
)
inputs["surface"][:, 1:3] = inputs["surface"][:, 1:3].astype(int)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is probably handled somewhere in the InputReader (and thus not related to this PR) but I couldn't find it at the first glance, do we ever check if the given l modes are all 0s? Even if the user gives nonzero values, they will probably be ignored but still, a warning or something could be good.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

They are not, I checked and inputs["surface"] has floats for the mode number parts of it. Probably we should handle in input reader, you're right

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

equilibrium.utils.parse_surface ensures that the surface array either has all l==0 or all n==0:

def parse_surface(surface, NFP=1, sym=True, spectral_indexing="ansi"):

# remove the keys (pertaining to continuation and solver tols)
# that an Equilibrium does not need
unused_keys = [
"pres_ratio",
"bdry_ratio",
"pert_order",
"ftol",
"xtol",
"gtol",
"maxiter",
"objective",
"optimizer",
"bdry_mode",
"output_path",
"verbose",
]
[inputs.pop(key) for key in unused_keys]
inputs.update(kwargs)
return cls(**inputs)

def solve(
self,
objective="force",
Expand Down
10 changes: 7 additions & 3 deletions desc/geometry/surface.py
Original file line number Diff line number Diff line change
Expand Up @@ -320,9 +320,13 @@ def from_input_file(cls, path, **kwargs):
inputs = InputReader().parse_inputs(path)[-1]
if (inputs["bdry_ratio"] is not None) and (inputs["bdry_ratio"] != 1):
warnings.warn(
"boundary_ratio = {} != 1, surface may not be as expected".format(
inputs["bdry_ratio"]
)
"`bdry_ratio` is intended as an input for the continuation method."
"`bdry_ratio`=1 uses the given surface modes as is, any other "
"scalar value will scale the non-axisymmetric modes by that "
"value. The final value of `bdry_ratio` in the input file is "
f"{inputs['bdry_ratio']}, this means the created "
"FourierRZToroidalSurface will be a scaled version of the "
"input file boundary."
)
surf = cls(
inputs["surface"][:, 3],
Expand Down
4 changes: 2 additions & 2 deletions desc/input_reader.py
Original file line number Diff line number Diff line change
Expand Up @@ -695,7 +695,7 @@ def write_desc_input(filename, inputs, header=""): # noqa: C901 - fxn too compl
f.write("# global parameters\n")
f.write("sym = {:1d} \n".format(inputs[0]["sym"]))
f.write("NFP = {:3d} \n".format(int(inputs[0]["NFP"])))
f.write("Psi = {:.8f} \n".format(inputs[0]["Psi"]))
f.write("Psi = {:.8e} \n".format(inputs[0]["Psi"]))

f.write("\n# spectral resolution\n")
for key, val in {
Expand Down Expand Up @@ -796,7 +796,7 @@ def desc_output_to_input( # noqa: C901 - fxn too complex
xtol=1e-6,
gtol=1e-6,
maxiter=100,
threshold=1e-10,
threshold=0,
):
"""Generate a DESC input file from a DESC output file.

Expand Down
6 changes: 6 additions & 0 deletions docs/api_equilibrium.rst
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,10 @@ An ``EquilibriaFamily`` is a ``list`` like object for storing multiple equilibri
desc.equilibrium.Equilibrium
desc.equilibrium.EquilibriaFamily

The ``Equilibrium`` class may be instantiated in a couple of ways in addition to providing inputs to its constructor.
- from an existing DESC or VMEC input file with its ``from_input_file`` method
- from a ``pyQSC`` ``Qsc`` or ``pyQIC`` ``Qic`` near-axis equilibrium with the ``Equilibrium``'s ``from_near_axis`` method


Geometry
********
Expand All @@ -35,6 +39,8 @@ the magnetic axis, cross section, and various space curves.
desc.geometry.SplineXYZCurve
desc.geometry.ZernikeRZToroidalSection

The ``FourierRZToroidalSurface`` and the ``FourierRZCurve`` classes may be instantiated from an existing DESC or VMEC input file with their ``from_input_file`` method.


Profiles
********
Expand Down
747 changes: 370 additions & 377 deletions docs/notebooks/tutorials/basic_equilibrium.ipynb

Large diffs are not rendered by default.

5 changes: 3 additions & 2 deletions tests/inputs/input.DSHAPE_desc
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
# This DESC input file was auto generated from the VMEC input file
# /DESC/tests/inputs/input.DSHAPE
# .//tests//inputs//input.DSHAPE
# on 10/30/2024 at 14:22:23.
# For details on the various options see https://desc-docs.readthedocs.io/en/stable/input.html

# global parameters
sym = 1
NFP = 1
Psi = 1.00000000
Psi = 1.00000000e+00

# spectral resolution
L_rad = 11
Expand Down
72 changes: 69 additions & 3 deletions tests/test_input_output.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,10 +57,10 @@ def test_vmec_input(tmpdir_factory):
lines_direct = f.readlines()
with open(path_converted_file) as f:
lines_converted = f.readlines()
# skip first 3 lines as they have date and pwd info
for line1, line2 in zip(lines_correct[3:], lines_converted[4:]):
# skip first 4 lines as they have date and pwd info
for line1, line2 in zip(lines_correct[4:], lines_converted[4:]):
assert line1.strip() == line2.strip()
for line1, line2 in zip(lines_correct[3:], lines_direct):
for line1, line2 in zip(lines_correct[4:], lines_direct):
assert line1.strip() == line2.strip()


Expand Down Expand Up @@ -226,6 +226,72 @@ def test_near_axis_input_files():
os.remove(".//tests//inputs//input.QSC_r2_5.5_vmec_desc")


@pytest.mark.unit
def test_from_input_file_equilibrium_desc_vmec_DSHAPE():
"""Test that from_input_file works for DESC input files."""
vmec_path = ".//tests//inputs//input.DSHAPE"
desc_path = ".//tests//inputs//input.DSHAPE_desc"
kwargs = {"spectral_indexing": "fringe"}
with pytest.warns(UserWarning, match="Left handed"):
eq = Equilibrium.from_input_file(desc_path, **kwargs)
with pytest.warns(UserWarning):
eq_VMEC = Equilibrium.from_input_file(vmec_path, **kwargs)

# make sure the loaded eqs are equivalent
np.testing.assert_allclose(eq.R_lmn, eq_VMEC.R_lmn)
np.testing.assert_allclose(eq.Z_lmn, eq_VMEC.Z_lmn)
np.testing.assert_allclose(eq.L_lmn, eq_VMEC.L_lmn)
np.testing.assert_allclose(eq.Rb_lmn, eq_VMEC.Rb_lmn)
np.testing.assert_allclose(eq.Zb_lmn, eq_VMEC.Zb_lmn)
np.testing.assert_allclose(eq.Ra_n, eq_VMEC.Ra_n)
np.testing.assert_allclose(eq.Za_n, eq_VMEC.Za_n)
np.testing.assert_allclose(eq.Psi, eq_VMEC.Psi)
assert eq.pressure.equiv(eq_VMEC.pressure)
assert eq.iota.equiv(eq_VMEC.iota)
assert eq.current is None
assert eq_VMEC.current is None
assert eq.sym == eq_VMEC.sym

# check against the DSHAPE bdry and profiles
eq_example = desc.examples.get("DSHAPE")
eq.change_resolution(L=eq_example.L, M=eq_example.M)
np.testing.assert_allclose(eq.Rb_lmn, eq_example.Rb_lmn)
np.testing.assert_allclose(eq.Zb_lmn, eq_example.Zb_lmn)
np.testing.assert_allclose(eq.Psi, eq_example.Psi)
np.testing.assert_allclose(eq.p_l, eq_example.p_l)
# our example's iota is negative of this input files's
np.testing.assert_allclose(eq.i_l, -eq_example.i_l)
assert eq.sym == eq_example.sym


@pytest.mark.unit
def test_from_input_file_equilibrium_desc_vmec():
"""Test that from_input_file gives same eq for DESC and VMEC input files."""
vmec_path = ".//tests//inputs//input.QSC_r2_5.5_vmec"
desc_path = ".//tests//inputs//input.QSC_r2_5.5_desc"
kwargs = {"L": 10, "M": 10, "N": 14}
eq = Equilibrium.from_input_file(desc_path, **kwargs)
with pytest.warns(UserWarning):
eq_VMEC = Equilibrium.from_input_file(vmec_path, **kwargs)

np.testing.assert_allclose(eq.R_lmn, eq_VMEC.R_lmn)
np.testing.assert_allclose(eq.Z_lmn, eq_VMEC.Z_lmn)
np.testing.assert_allclose(eq.L_lmn, eq_VMEC.L_lmn)
np.testing.assert_allclose(eq.Rb_lmn, eq_VMEC.Rb_lmn)
np.testing.assert_allclose(eq.Zb_lmn, eq_VMEC.Zb_lmn)
np.testing.assert_allclose(eq.Ra_n, eq_VMEC.Ra_n)
np.testing.assert_allclose(eq.Za_n, eq_VMEC.Za_n)
np.testing.assert_allclose(eq.Psi, eq_VMEC.Psi)
assert eq.pressure.equiv(eq_VMEC.pressure)
assert eq.current.equiv(eq_VMEC.current)
assert eq.iota is None
assert eq_VMEC.iota is None
assert eq.sym == eq_VMEC.sym

if os.path.exists(".//tests//inputs//input.QSC_r2_5.5_vmec_desc"):
os.remove(".//tests//inputs//input.QSC_r2_5.5_vmec_desc")


@pytest.mark.unit
def test_vmec_input_surface_threshold():
"""Test ."""
Expand Down