Skip to content

Commit

Permalink
fixed suggestions on this summers' PR
Browse files Browse the repository at this point in the history
  • Loading branch information
smiet committed Nov 12, 2024
1 parent 57c8291 commit c52dd99
Showing 1 changed file with 23 additions and 15 deletions.
38 changes: 23 additions & 15 deletions src/simsopt/configs/zoo.py
Original file line number Diff line number Diff line change
Expand Up @@ -198,21 +198,27 @@ def get_w7x_data(Nt_coils=48, Nt_ma=10, ppp=2):


@SimsoptRequires(requests is not None, "You need to install the requests library to use this function. Run 'pip install requests'")
def get_QUASR_data(ID, return_style='default'):
def get_QUASR_data(ID, return_style='quasr-style'):
"""
Download a configuration from the QUASAR database.
Download a configuration from the QUASR database.
Args:
ID: the ID of the configuration to download. The database is navigatable at https://quasr.flatironinstitute.org/
Alternatively, you can download the latest full set of devices from https://zenodo.org/doi/10.5281/zenodo.10050655
return_style: 'default' or 'json'. If 'default', the function will return the curves, currents and magnetic axis
like the other configurations in the zoo.
If 'json' the function will return the full set of coils, magnetic axis and a number of surfaces
parametrized in Boozer coordinates.
return_style: 'simsopt-style' or 'quasr-style'. '.
simsopt-style: [coils_1fp, surface], similar to get_ncsx_data(), gives the curves and currenst for one field period,
which you can copy and rotate using simsopt methods (allows finer control over degrees-of-freedom).
NOTE: this does not return the magnetic axis.
quasr-style: [coils_all, surfaces], returns all coils and all surfaces in the object.
returns: depending on return_style:
simsopt-style: [list of simsopt.geo.Coil objects, list of simsopt.field.Current objects]
quasr-style: [list of simsopt.geo.SurfaceXYZTensorFourier objects, list of simsopt.field.COIL objects]
"""

assert return_style in ['default', 'json']
if return_style not in ['simsopt-style', 'quasr-style']:
raise ValueError(f"invalid return_style: {return_style}, must be either simsopt-style or quasr-style")

Check warning on line 221 in src/simsopt/configs/zoo.py

View check run for this annotation

Codecov / codecov/patch

src/simsopt/configs/zoo.py#L220-L221

Added lines #L220 - L221 were not covered by tests

id_str = f"{ID:07d}"

Check warning on line 223 in src/simsopt/configs/zoo.py

View check run for this annotation

Codecov / codecov/patch

src/simsopt/configs/zoo.py#L223

Added line #L223 was not covered by tests
# string to 7 digits
Expand All @@ -221,18 +227,20 @@ def get_QUASR_data(ID, return_style='default'):
with requests.get(url) as r:
if r.status_code == 200:
print(f"Configuration with ID {ID:07} downloaded successfully")
surfaces, ma, coils = json.loads(r.content, cls=GSONDecoder)
surfaces, coils = json.loads(r.content, cls=GSONDecoder)

Check warning on line 230 in src/simsopt/configs/zoo.py

View check run for this annotation

Codecov / codecov/patch

src/simsopt/configs/zoo.py#L227-L230

Added lines #L227 - L230 were not covered by tests
else:
raise ValueError(f"Configuration ID {ID:07d} does not exist. Status code: {r.status_code}")
raise ValueError(f"Download of ID {ID:07d} failed. Status code: {r.status_code}\n Check if the confituration exists")

Check warning on line 232 in src/simsopt/configs/zoo.py

View check run for this annotation

Codecov / codecov/patch

src/simsopt/configs/zoo.py#L232

Added line #L232 was not covered by tests

if return_style == 'default':
nfp = ma.nfp
nc_per_hp = len(coils) // nfp // (1 + ma.stellsym)
if return_style == 'simsopt-style':
nfp = surfaces[0].nfp
nc_per_hp = len(coils) // nfp // (1 + surfaces[0].stellsym)
coils = coils[:nc_per_hp]

Check warning on line 237 in src/simsopt/configs/zoo.py

View check run for this annotation

Codecov / codecov/patch

src/simsopt/configs/zoo.py#L234-L237

Added lines #L234 - L237 were not covered by tests

curves = [coil.curve for coil in coils]
currents = [coil.current for coil in coils]
return curves, currents, ma
return curves, currents
elif return_style == 'quasr-style':
return surfaces, coils

Check warning on line 243 in src/simsopt/configs/zoo.py

View check run for this annotation

Codecov / codecov/patch

src/simsopt/configs/zoo.py#L239-L243

Added lines #L239 - L243 were not covered by tests
else:
raise ValueError #should not be reached as we check before download to avoid clobbering the database.

Check warning on line 245 in src/simsopt/configs/zoo.py

View check run for this annotation

Codecov / codecov/patch

src/simsopt/configs/zoo.py#L245

Added line #L245 was not covered by tests

elif return_style == 'json':
return coils, ma, surfaces

0 comments on commit c52dd99

Please sign in to comment.