constellaration 0.2.6__tar.gz → 0.3.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (58) hide show
  1. {constellaration-0.2.6 → constellaration-0.3.0}/.gitignore +6 -0
  2. {constellaration-0.2.6 → constellaration-0.3.0}/LICENSE +1 -1
  3. {constellaration-0.2.6 → constellaration-0.3.0}/PKG-INFO +53 -10
  4. {constellaration-0.2.6 → constellaration-0.3.0}/README.md +46 -3
  5. constellaration-0.3.0/constellaration/__about__.py +24 -0
  6. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/boozer/boozer.py +9 -9
  7. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/data_generation/vmec_optimization.py +1 -1
  8. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/data_generation/vmec_optimization_settings.py +2 -1
  9. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/generative_model/bootstrap_dataset.py +4 -4
  10. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/generative_model/mcmc_algo.py +3 -3
  11. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/generative_model/optimize_with_mcmc.py +6 -5
  12. constellaration-0.3.0/constellaration/geometry/spectral_condensation.py +356 -0
  13. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/geometry/surface_rz_fourier.py +329 -25
  14. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/geometry/surface_utils.py +1 -1
  15. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/geometry/surface_utils_desc.py +9 -14
  16. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/initial_guess.py +1 -1
  17. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/mhd/geometry_utils.py +3 -3
  18. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/mhd/magnetics_utils.py +1 -1
  19. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/mhd/near_axis_configuration.py +1 -1
  20. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/mhd/turbulent_transport.py +4 -2
  21. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/mhd/vmec_settings.py +9 -13
  22. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/mhd/vmec_utils.py +19 -7
  23. constellaration-0.3.0/constellaration/mhd_stable_qi_scoring.py +445 -0
  24. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/omnigeneity/omnigenity_field.py +24 -24
  25. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/omnigeneity/omnigenity_field_desc.py +2 -2
  26. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/omnigeneity/omnigenity_field_plot.py +4 -4
  27. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/omnigeneity/omnigenity_field_sampling.py +7 -7
  28. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/optimization/augmented_lagrangian.py +3 -11
  29. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/problems.py +48 -4
  30. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/utils/pytree.py +13 -10
  31. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/utils/visualization.py +1 -1
  32. {constellaration-0.2.6 → constellaration-0.3.0}/pyproject.toml +12 -6
  33. constellaration-0.2.6/constellaration/__about__.py +0 -34
  34. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/__init__.py +0 -0
  35. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/boozer/__init__.py +0 -0
  36. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/data_generation/__init__.py +0 -0
  37. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/data_generation/desc_optimization.py +0 -0
  38. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/data_generation/desc_optimization_settings.py +0 -0
  39. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/data_generation/rotational_transform_scaling.py +0 -0
  40. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/forward_model.py +0 -0
  41. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/geometry/__init__.py +0 -0
  42. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/geometry/radial_profile.py +0 -0
  43. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/mhd/__init__.py +0 -0
  44. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/mhd/flux_power_series.py +0 -0
  45. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/mhd/ideal_mhd_parameters.py +0 -0
  46. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/mhd/ideal_mhd_parameters_utils.py +0 -0
  47. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/omnigeneity/__init__.py +0 -0
  48. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/omnigeneity/qi.py +0 -0
  49. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/optimization/__init__.py +0 -0
  50. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/optimization/augmented_lagrangian_runner.py +0 -0
  51. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/optimization/scipy_minimize_runner.py +0 -0
  52. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/optimization/settings.py +0 -0
  53. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/utils/__init__.py +0 -0
  54. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/utils/dataset_utils.py +0 -0
  55. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/utils/file_exporter.py +0 -0
  56. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/utils/seed_util.py +0 -0
  57. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/utils/types.py +0 -0
  58. {constellaration-0.2.6 → constellaration-0.3.0}/constellaration/utils/visualization_utils.py +0 -0
@@ -172,3 +172,9 @@ cython_debug/
172
172
 
173
173
  # PyPI configuration file
174
174
  .pypirc
175
+
176
+ # Presentation visualizations
177
+ presentations/
178
+
179
+ # Generated by hatch-vcs at build time from the git tag
180
+ src/constellaration/__about__.py
@@ -1,6 +1,6 @@
1
1
  MIT License
2
2
 
3
- Copyright (c) 2025-present Proxima Fusion Gmbh <info@proximafusion.com>
3
+ Copyright (c) 2025-present Proxima Fusion GmbH <info@proximafusion.com>
4
4
 
5
5
  Permission is hereby granted, free of charge, to any person obtaining a copy
6
6
  of this software and associated documentation files (the "Software"), to deal
@@ -1,6 +1,6 @@
1
- Metadata-Version: 2.4
1
+ Metadata-Version: 2.5
2
2
  Name: constellaration
3
- Version: 0.2.6
3
+ Version: 0.3.0
4
4
  Summary: Code for analyzing and evaluating stellarator plasma shapes
5
5
  Project-URL: Documentation, https://github.com/proximafusion/constellaration#readme
6
6
  Project-URL: Issues, https://github.com/proximafusion/constellaration/issues
@@ -18,17 +18,17 @@ Classifier: Programming Language :: Python :: Implementation :: PyPy
18
18
  Requires-Python: >=3.10
19
19
  Requires-Dist: booz-xform==0.0.9
20
20
  Requires-Dist: datasets==3.5.1
21
- Requires-Dist: desc-opt
22
- Requires-Dist: interpax==0.3.8
21
+ Requires-Dist: desc-opt==0.17.1
22
+ Requires-Dist: interpax==0.3.14
23
23
  Requires-Dist: ipykernel
24
24
  Requires-Dist: ipython
25
- Requires-Dist: jax<=0.6.0
25
+ Requires-Dist: jax<=0.10.0
26
26
  Requires-Dist: jaxtyping
27
27
  Requires-Dist: kaleido
28
28
  Requires-Dist: matplotlib==3.9.2
29
29
  Requires-Dist: nbformat
30
30
  Requires-Dist: nevergrad
31
- Requires-Dist: numpy>=2.1.3
31
+ Requires-Dist: numpy<2.3,>=2.1.3
32
32
  Requires-Dist: orjson==3.10.16
33
33
  Requires-Dist: pandas==2.2.3
34
34
  Requires-Dist: plotly<7.0.0
@@ -43,7 +43,7 @@ Requires-Dist: scipy
43
43
  Requires-Dist: seaborn==0.13.2
44
44
  Requires-Dist: simsopt==1.8.1; python_version < '3.13'
45
45
  Requires-Dist: simsopt==1.9.0; python_version >= '3.13'
46
- Requires-Dist: vmecpp==0.4.11
46
+ Requires-Dist: vmecpp==0.5.4
47
47
  Provides-Extra: lint
48
48
  Requires-Dist: pre-commit; extra == 'lint'
49
49
  Provides-Extra: test
@@ -58,10 +58,45 @@ Description-Content-Type: text/markdown
58
58
 
59
59
  # ConStellaration: A dataset of QI-like stellarator plasma boundaries and optimization benchmarks
60
60
 
61
- [ConStellaration](https://arxiv.org/abs/2506.19583) is a dataset of diverse QI-like stellarator plasma boundary shapes and optimization benchmakrs, paired with their ideal-MHD equilibria and performance metrics.
61
+ [ConStellaration](https://arxiv.org/abs/2506.19583) is a dataset of diverse QI-like stellarator plasma boundary shapes and optimization benchmarks, paired with their ideal-MHD equilibria and performance metrics.
62
62
  The dataset is available on [Hugging Face](https://huggingface.co/datasets/proxima-fusion/constellaration).
63
63
  The repository contains a suite of tools and notebooks for exploring the dataset, including a forward model for plasma simulation, scoring functions for optimization evaluation and data-driven generative modeling.
64
64
 
65
+ ## Reproducibility and benchmark versions
66
+
67
+ This repository is under active development: new metrics, problems and dependency updates (e.g. [VMEC++](https://github.com/proximafusion/vmecpp)) can change the values computed by the forward model and the scoring functions.
68
+ To keep results comparable, the benchmark of [Cadena et al., NeurIPS 2025](https://openreview.net/forum?id=NQSbGKlCpx) is frozen to a tagged release.
69
+
70
+ | Use case | Version |
71
+ | --- | --- |
72
+ | [ConStellaration leaderboard](https://huggingface.co/spaces/proxima-fusion/constellaration-bench) (geometrical, simple-to-build QI, MHD-stable QI problems) | [`v0.3.0`](https://github.com/proximafusion/constellaration/releases/tag/v0.3.0) |
73
+ | Original NeurIPS 2025 paper experiments | [`0.2.1`](https://pypi.org/project/constellaration/0.2.1/) |
74
+ | Latest development | `main` |
75
+
76
+ If you are working on the challenge or comparing against the leaderboard, install the pinned benchmark release:
77
+
78
+ ```bash
79
+ pip install constellaration==0.3.0
80
+ ```
81
+
82
+ or, from source:
83
+
84
+ ```bash
85
+ git clone --branch v0.3.0 https://github.com/proximafusion/constellaration.git
86
+ cd constellaration
87
+ pip install .
88
+ ```
89
+
90
+ The leaderboard evaluates every submission with exactly this version, so scores computed locally with `constellaration==0.3.0` match the leaderboard.
91
+ Releases newer than `v0.3.0` may produce different metrics and scores and are not used by the leaderboard until explicitly announced here.
92
+ The PyPI version always equals the git tag of the release it was built from (`constellaration==X.Y.Z` ⇔ tag `vX.Y.Z`).
93
+
94
+ Differences of `v0.3.0` with respect to `0.2.1` (used for the original paper experiments) that can affect the benchmark scores:
95
+ - The rotational transform constraint uses the absolute value of the edge rotational transform, i.e. boundaries with negative iota are no longer penalized.
96
+ - Newer versions of VMEC++ and of other numerical dependencies (`booz-xform`, `jax`, `desc-opt`) are used, which can cause small numerical differences.
97
+
98
+ To reproduce the exact numbers of the paper, use `pip install constellaration==0.2.1` (this version is no longer maintained and may require pinning older versions of its unpinned dependencies).
99
+
65
100
  ## Installation
66
101
 
67
102
  The following instructions have been tested on **Ubuntu 22.04** and **Ubuntu 24.04**. Other platforms may require additional steps and have not been validated.
@@ -92,7 +127,7 @@ pip install constellaration
92
127
 
93
128
  2. Install the required system dependencies
94
129
  1. **On Ubuntu**: `sudo apt-get update && sudo apt-get install -y libnetcdf-dev`
95
- 2. **On MAC-OS**: `brew install netcdf`
130
+ 2. **On macOS**: `brew install netcdf`
96
131
 
97
132
  3. Install the required Python dependencies:
98
133
 
@@ -100,6 +135,14 @@ pip install constellaration
100
135
  pip install .
101
136
  ```
102
137
 
138
+ **Note for macOS:** building `booz-xform` from source calls `python` from the `PATH`. If `python` does not resolve to the interpreter you are installing into (e.g. with a pyenv `system` global, where only `python3` exists), the build fails with `Could not find a package configuration file provided by "pybind11"`. Install into an activated virtual environment so that `python` points to it:
139
+
140
+ ```bash
141
+ python3 -m venv .venv
142
+ source .venv/bin/activate
143
+ pip install .
144
+ ```
145
+
103
146
  ### Running with Docker
104
147
 
105
148
  If you prefer not to install system dependencies, you can use the provided Dockerfile to build a Docker image and run your scripts in a container.
@@ -135,7 +178,7 @@ pip install -e ".[test,lint]"
135
178
 
136
179
  We use **pre-commit** to automatically lint and format code before each commit. Linting is static code analysis that catches style issues and potential errors. If any **hook** fails, the commit will be blocked until you fix the reported issues and re-stage your changes.
137
180
 
138
- Install the hook (once per clone):
181
+ Install the hook (once per clone):
139
182
  ```bash
140
183
  pip install pre-commit
141
184
  pre-commit install
@@ -6,10 +6,45 @@
6
6
 
7
7
  # ConStellaration: A dataset of QI-like stellarator plasma boundaries and optimization benchmarks
8
8
 
9
- [ConStellaration](https://arxiv.org/abs/2506.19583) is a dataset of diverse QI-like stellarator plasma boundary shapes and optimization benchmakrs, paired with their ideal-MHD equilibria and performance metrics.
9
+ [ConStellaration](https://arxiv.org/abs/2506.19583) is a dataset of diverse QI-like stellarator plasma boundary shapes and optimization benchmarks, paired with their ideal-MHD equilibria and performance metrics.
10
10
  The dataset is available on [Hugging Face](https://huggingface.co/datasets/proxima-fusion/constellaration).
11
11
  The repository contains a suite of tools and notebooks for exploring the dataset, including a forward model for plasma simulation, scoring functions for optimization evaluation and data-driven generative modeling.
12
12
 
13
+ ## Reproducibility and benchmark versions
14
+
15
+ This repository is under active development: new metrics, problems and dependency updates (e.g. [VMEC++](https://github.com/proximafusion/vmecpp)) can change the values computed by the forward model and the scoring functions.
16
+ To keep results comparable, the benchmark of [Cadena et al., NeurIPS 2025](https://openreview.net/forum?id=NQSbGKlCpx) is frozen to a tagged release.
17
+
18
+ | Use case | Version |
19
+ | --- | --- |
20
+ | [ConStellaration leaderboard](https://huggingface.co/spaces/proxima-fusion/constellaration-bench) (geometrical, simple-to-build QI, MHD-stable QI problems) | [`v0.3.0`](https://github.com/proximafusion/constellaration/releases/tag/v0.3.0) |
21
+ | Original NeurIPS 2025 paper experiments | [`0.2.1`](https://pypi.org/project/constellaration/0.2.1/) |
22
+ | Latest development | `main` |
23
+
24
+ If you are working on the challenge or comparing against the leaderboard, install the pinned benchmark release:
25
+
26
+ ```bash
27
+ pip install constellaration==0.3.0
28
+ ```
29
+
30
+ or, from source:
31
+
32
+ ```bash
33
+ git clone --branch v0.3.0 https://github.com/proximafusion/constellaration.git
34
+ cd constellaration
35
+ pip install .
36
+ ```
37
+
38
+ The leaderboard evaluates every submission with exactly this version, so scores computed locally with `constellaration==0.3.0` match the leaderboard.
39
+ Releases newer than `v0.3.0` may produce different metrics and scores and are not used by the leaderboard until explicitly announced here.
40
+ The PyPI version always equals the git tag of the release it was built from (`constellaration==X.Y.Z` ⇔ tag `vX.Y.Z`).
41
+
42
+ Differences of `v0.3.0` with respect to `0.2.1` (used for the original paper experiments) that can affect the benchmark scores:
43
+ - The rotational transform constraint uses the absolute value of the edge rotational transform, i.e. boundaries with negative iota are no longer penalized.
44
+ - Newer versions of VMEC++ and of other numerical dependencies (`booz-xform`, `jax`, `desc-opt`) are used, which can cause small numerical differences.
45
+
46
+ To reproduce the exact numbers of the paper, use `pip install constellaration==0.2.1` (this version is no longer maintained and may require pinning older versions of its unpinned dependencies).
47
+
13
48
  ## Installation
14
49
 
15
50
  The following instructions have been tested on **Ubuntu 22.04** and **Ubuntu 24.04**. Other platforms may require additional steps and have not been validated.
@@ -40,7 +75,7 @@ pip install constellaration
40
75
 
41
76
  2. Install the required system dependencies
42
77
  1. **On Ubuntu**: `sudo apt-get update && sudo apt-get install -y libnetcdf-dev`
43
- 2. **On MAC-OS**: `brew install netcdf`
78
+ 2. **On macOS**: `brew install netcdf`
44
79
 
45
80
  3. Install the required Python dependencies:
46
81
 
@@ -48,6 +83,14 @@ pip install constellaration
48
83
  pip install .
49
84
  ```
50
85
 
86
+ **Note for macOS:** building `booz-xform` from source calls `python` from the `PATH`. If `python` does not resolve to the interpreter you are installing into (e.g. with a pyenv `system` global, where only `python3` exists), the build fails with `Could not find a package configuration file provided by "pybind11"`. Install into an activated virtual environment so that `python` points to it:
87
+
88
+ ```bash
89
+ python3 -m venv .venv
90
+ source .venv/bin/activate
91
+ pip install .
92
+ ```
93
+
51
94
  ### Running with Docker
52
95
 
53
96
  If you prefer not to install system dependencies, you can use the provided Dockerfile to build a Docker image and run your scripts in a container.
@@ -83,7 +126,7 @@ pip install -e ".[test,lint]"
83
126
 
84
127
  We use **pre-commit** to automatically lint and format code before each commit. Linting is static code analysis that catches style issues and potential errors. If any **hook** fails, the commit will be blocked until you fix the reported issues and re-stage your changes.
85
128
 
86
- Install the hook (once per clone):
129
+ Install the hook (once per clone):
87
130
  ```bash
88
131
  pip install pre-commit
89
132
  pre-commit install
@@ -0,0 +1,24 @@
1
+ # file generated by vcs-versioning
2
+ # don't change, don't track in version control
3
+ from __future__ import annotations
4
+
5
+ __all__ = [
6
+ "__version__",
7
+ "__version_tuple__",
8
+ "version",
9
+ "version_tuple",
10
+ "__commit_id__",
11
+ "commit_id",
12
+ ]
13
+
14
+ version: str
15
+ __version__: str
16
+ __version_tuple__: tuple[int | str, ...]
17
+ version_tuple: tuple[int | str, ...]
18
+ commit_id: str | None
19
+ __commit_id__: str | None
20
+
21
+ __version__ = version = '0.3.0'
22
+ __version_tuple__ = version_tuple = (0, 3, 0)
23
+
24
+ __commit_id__ = commit_id = None
@@ -41,7 +41,7 @@ class BoozerOutput(pydantic.BaseModel, arbitrary_types_allowed=True):
41
41
  """The number of Boozer poloidal modes."""
42
42
 
43
43
  max_boozer_toroidal_mode: int
44
- """The number of Boozer toroidal modes."""
44
+ """The maximum Boozer toroidal mode number."""
45
45
 
46
46
  n_poloidal_modes: int
47
47
  """The number of poloidal modes."""
@@ -100,10 +100,10 @@ class BoozerOutput(pydantic.BaseModel, arbitrary_types_allowed=True):
100
100
  """The magnetic field strength sine Fourier coefficients in Boozer coordinates."""
101
101
 
102
102
  gmnc_b: FOURIER_COEFFICIENTS
103
- """The equilibrium jacobian cosine Fourier coefficients in Boozer coordinates."""
103
+ """The equilibrium Jacobian cosine Fourier coefficients in Boozer coordinates."""
104
104
 
105
105
  gmns_b: FOURIER_COEFFICIENTS
106
- """The equilibrium jacobian sine Fourier coefficients in Boozer coordinates."""
106
+ """The equilibrium Jacobian sine Fourier coefficients in Boozer coordinates."""
107
107
 
108
108
  rmnc_b: FOURIER_COEFFICIENTS
109
109
  """The equilibrium flux surfaces R coordinate cosine Fourier coefficients in Boozer
@@ -216,7 +216,7 @@ class BoozerSettings(pydantic.BaseModel):
216
216
  """
217
217
 
218
218
  verbose: bool = False
219
- """If True, running booz_xform in verbose mode."""
219
+ """If True, runs booz_xform in verbose mode."""
220
220
 
221
221
 
222
222
  class BoozerPresetSettings(pydantic.BaseModel):
@@ -230,7 +230,7 @@ class BoozerPresetSettings(pydantic.BaseModel):
230
230
  """
231
231
 
232
232
  verbose: bool = False
233
- """If True, running booz_xform in verbose mode."""
233
+ """If True, runs booz_xform in verbose mode."""
234
234
 
235
235
 
236
236
  def run_boozer(
@@ -305,10 +305,10 @@ def create_boozer_settings_from_equilibrium_resolution(
305
305
  mhd_equilibrium: vmec_utils.VmecppWOut,
306
306
  settings: BoozerPresetSettings,
307
307
  ) -> BoozerSettings:
308
- """Derives Boozer transformation settings from the resolution of the equilibriumn
309
- according to the given boozer settings preset.
308
+ """Derives Boozer transformation settings from the resolution of the equilibrium
309
+ according to the given Boozer settings preset.
310
310
 
311
- See the wrapped function boozer.boozer_settings_from_equilibrium_resolution for
311
+ See the wrapped function boozer.boozer_settings_from_equilibrium_resolution for
312
312
  more details.
313
313
  """
314
314
  return boozer_settings_from_equilibrium_resolution(
@@ -333,7 +333,7 @@ def boozer_settings_from_equilibrium_resolution(
333
333
  normalized_toroidal_flux: The normalized toroidal flux values at which the
334
334
  transformation to Boozer coordinates is performed. If None, the
335
335
  transformation to Boozer coordinates will be run on all flux surfaces.
336
- verbose: If True, running booz_xform in verbose mode.
336
+ verbose: If True, runs booz_xform in verbose mode.
337
337
  """
338
338
  mpol = max(6 * mhd_equilibrium.mpol, 2)
339
339
  ntor = max(2 * mhd_equilibrium.ntor - 1, 0)
@@ -273,7 +273,7 @@ def _compute_omnigenous_field_residuals(
273
273
  modB: jt.Float[np.ndarray, " n_collocation_points"],
274
274
  target_modB: jt.Float[np.ndarray, " n_collocation_points"],
275
275
  ) -> jt.Float[np.ndarray, " n_collocation_points"]:
276
- # Scale both magnetic field strength such that the average is ~1T.
276
+ # Scale both magnetic field strengths such that the average is ~1T.
277
277
  scaled_modB = modB / np.mean(modB)
278
278
  scaled_target_modB = target_modB / np.mean(target_modB)
279
279
  return (scaled_modB - scaled_target_modB) / scaled_target_modB
@@ -12,7 +12,8 @@ class OmnigenousFieldVmecOptimizationSettings(pydantic.BaseModel):
12
12
  """The maximum toroidal mode number of the resulting configuration."""
13
13
 
14
14
  n_inner_optimizations: int = 1
15
- """The number of inner optimization to perform to obtain the final configuration."""
15
+ """The number of inner optimizations to perform to obtain the final
16
+ configuration."""
16
17
 
17
18
  gradient_free_budget_per_design_variable: int = 100
18
19
  """The budget for the gradient free optimization per design variable.
@@ -19,10 +19,10 @@ ANY_PROBLEM = problems.GeometricalProblem | problems.SimpleToBuildQIStellarator
19
19
 
20
20
 
21
21
  def load_source_datasets_with_no_errors() -> pd.DataFrame:
22
- """Load and concatenate source datasets from HugginFace dataset.
22
+ """Load and concatenate source datasets from the Hugging Face dataset.
23
23
 
24
24
  Includes all ids, unfolded targets and unfolded metrics but not the contents of the
25
- mhd equilibria.
25
+ MHD equilibria.
26
26
  """
27
27
  dframe = datasets.load_dataset("proxima-fusion/constellaration", "default")[ # type: ignore
28
28
  "train"
@@ -213,7 +213,7 @@ def _to_Y_constraints(
213
213
  dframe: pd.DataFrame,
214
214
  problem: ANY_PROBLEM,
215
215
  ) -> np.ndarray:
216
- """Extracts the constraint metrics from the DataFrame from the given problem."""
216
+ """Extracts the constraint metrics from the DataFrame for the given problem."""
217
217
  if isinstance(problem, problems.GeometricalProblem):
218
218
  targets = np.array(
219
219
  [
@@ -264,7 +264,7 @@ def _to_Y_objective(
264
264
  dframe: pd.DataFrame,
265
265
  problem: ANY_PROBLEM,
266
266
  ) -> np.ndarray:
267
- """Extracts the objective metrics from the DataFrame from the given problem."""
267
+ """Extracts the objective metrics from the DataFrame for the given problem."""
268
268
  if isinstance(problem, problems.GeometricalProblem):
269
269
  return dframe["max_elongation"].to_numpy()
270
270
  elif isinstance(problem, problems.SimpleToBuildQIStellarator):
@@ -108,9 +108,9 @@ class RandomWalkMetropolis:
108
108
  def _tune_scale_covariance(
109
109
  self, proposal_scale: float, accept_rate: float
110
110
  ) -> float:
111
- r"""Tune the acceptance rate according to the last tuning interval. If higher
112
- acceptance rate , means you need to expand your search field or increase
113
- variance(its too small currently)
111
+ r"""Tune the acceptance rate according to the last tuning interval. If the
112
+ acceptance rate is higher, you need to expand your search field or increase
113
+ the variance (it's too small currently).
114
114
 
115
115
  The goal is an acceptance rate within 20\% - 50\%.
116
116
  The (acceptance) rate is adapted according to the following rule:
@@ -30,7 +30,8 @@ def optimize_with_mcmc(
30
30
  initial_stepsize: np.ndarray | None = None,
31
31
  **kwargs,
32
32
  ) -> np.ndarray:
33
- """Optimizes a function using MCMC sampling."
33
+ """Optimizes a function using MCMC sampling.
34
+
34
35
  Args:
35
36
  function: For a given design variable, a function to compute a tuple of
36
37
  objectives and constraints. Constraint negative is considered as
@@ -68,13 +69,13 @@ def optimize_with_mcmc(
68
69
  # for the objective, quasi-likelihood
69
70
  beta = settings.beta # can be a hyperparameter too. e.g. 1.0
70
71
  num_objectives = objectives.shape[0]
71
- # The is just a basic form to deal with multi-objective optimization
72
+ # This is just a basic form to deal with multi-objective optimization
72
73
  w = 1 / num_objectives
73
74
  quasi_likelihood_obj = np.exp(-beta * np.sum(w * objectives))
74
75
 
75
76
  # for the constraints, indicator function
76
77
  if settings.smooth_constraints:
77
- k = settings.k # more the k, more steep the sigmoid. default 3
78
+ k = settings.k # the larger k, the steeper the sigmoid. default 3
78
79
  sigma = (
79
80
  settings.sigma
80
81
  ) # controls the degree of constraint violation allowed, def 0.01
@@ -83,7 +84,7 @@ def optimize_with_mcmc(
83
84
  # sigmoid function
84
85
  quasi_likelihood_cons_tmp = 1.0 / (
85
86
  1.0 + np.exp(k * con)
86
- ) # 1/(1+exp(-k(-x))) -x as we want 1 when constraoints satisfied
87
+ ) # 1/(1+exp(-k(-x))) -x as we want 1 when constraints satisfied
87
88
  # else 0
88
89
  # gaussian
89
90
  quasi_likelihood_cons *= np.exp(
@@ -106,7 +107,7 @@ def optimize_with_mcmc(
106
107
 
107
108
  # Initialize the MCMC algorithm
108
109
  print(
109
- "Running MCMC (Adaptive Random Walk Metrolopis Hastings)"
110
+ "Running MCMC (Adaptive Random Walk Metropolis-Hastings) "
110
111
  f"with {settings.num_samples} samples"
111
112
  )
112
113
  mcmc = mcmc_algo.RandomWalkMetropolis(target_logprob=_log_target)