Debugging singular integrals #1440
Conversation
|
Might be related, but if you change the source grid NFP, you get different convergence/no convergence for the test for the axisymmetric case |
|
Maybe check basis being used? |
|
https://arxiv.org/pdf/2404.02799 |
Still seeing that the output of DFTInterpolator is not deterministic:
Here is repeated output of an example test that differs on different runs.
(desc-env-pip) kaya@fedora:~/Documents/project/DESC$ pytest -k test_singular_integral_vac_estell
/home/kaya/miniconda3/envs/desc-env-pip/lib/python3.12/multiprocessing/popen_fork.py:66: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.
self.pid = os.fork()
=================================================== test session starts ===================================================
platform linux -- Python 3.12.0, pytest-8.3.4, pluggy-1.5.0
Matplotlib: 3.9.3
Freetype: 2.6.1
benchmark: 4.0.0 (defaults: timer=time.perf_counter disable_gc=False min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=False warmup_iterations=100000)
rootdir: /home/kaya/Documents/project/DESC
configfile: setup.cfg
plugins: typeguard-2.13.3, jaxtyping-0.2.28, split-0.8.2, mpl-0.16.1, monitor-1.6.6, cov-5.0.0, benchmark-4.0.0, nbmake-1.5.3
collected 844 items / 843 deselected / 1 selected
tests/test_integrals.py F [100%]
======================================================== FAILURES =========================================================
___________________________________ TestSingularities.test_singular_integral_vac_estell ___________________________________
self = <tests.test_integrals.TestSingularities object at 0x7f6cf125db20>
@pytest.mark.unit
def test_singular_integral_vac_estell(self):
"""Test calculating Bplasma for vacuum estell, which should be near 0."""
eq = get("ESTELL")
grid = LinearGrid(M=25, N=25, NFP=eq.NFP)
keys = [
"K_vc",
"B",
"|B|^2",
"R",
"phi",
"Z",
"e^rho",
"n_rho",
"|e_theta x e_zeta|",
]
data = eq.compute(keys, grid=grid)
k = min(grid.num_theta, grid.num_zeta * grid.NFP)
s = k - 1
q = k // 2 + int(np.sqrt(k))
interpolator = DFTInterpolator(grid, grid, s, q)
Bplasma = virtual_casing_biot_savart(data, data, interpolator, loop=True)
# need extra factor of B/2 bc we're evaluating on plasma surface
Bplasma += data["B"] / 2
Bplasma = np.linalg.norm(Bplasma, axis=-1)
# scale by total field magnitude
B = Bplasma / np.linalg.norm(data["B"], axis=-1).mean()
> np.testing.assert_allclose(B, 0, atol=0.005)
tests/test_integrals.py:675:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
args = (<function assert_allclose.<locals>.compare at 0x7f6cf139d580>, array([0.26020428, 0.25978221, 0.25833294, ..., 0.25034399, 0.25539644,
0.25843631]), array(0))
kwds = {'equal_nan': True, 'err_msg': '', 'header': 'Not equal to tolerance rtol=1e-07, atol=0.005', 'verbose': True}
@wraps(func)
def inner(*args, **kwds):
with self._recreate_cm():
> return func(*args, **kwds)
E AssertionError:
E Not equal to tolerance rtol=1e-07, atol=0.005
E
E Mismatched elements: 2601 / 2601 (100%)
E Max absolute difference: 0.45300952
E Max relative difference: inf
E x: array([0.260204, 0.259782, 0.258333, ..., 0.250344, 0.255396, 0.258436])
E y: array(0)
../../../miniconda3/envs/desc-env-pip/lib/python3.12/contextlib.py:81: AssertionError
================================================= short test summary info =================================================
FAILED tests/test_integrals.py::TestSingularities::test_singular_integral_vac_estell - AssertionError:
=========================================== 1 failed, 843 deselected in 17.49s ============================================
(desc-env-pip) kaya@fedora:~/Documents/project/DESC$ pytest -k test_singular_integral_vac_estell
/home/kaya/miniconda3/envs/desc-env-pip/lib/python3.12/multiprocessing/popen_fork.py:66: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.
self.pid = os.fork()
=================================================== test session starts ===================================================
platform linux -- Python 3.12.0, pytest-8.3.4, pluggy-1.5.0
Matplotlib: 3.9.3
Freetype: 2.6.1
benchmark: 4.0.0 (defaults: timer=time.perf_counter disable_gc=False min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=False warmup_iterations=100000)
rootdir: /home/kaya/Documents/project/DESC
configfile: setup.cfg
plugins: typeguard-2.13.3, jaxtyping-0.2.28, split-0.8.2, mpl-0.16.1, monitor-1.6.6, cov-5.0.0, benchmark-4.0.0, nbmake-1.5.3
collected 844 items / 843 deselected / 1 selected
tests/test_integrals.py F [100%]
======================================================== FAILURES =========================================================
___________________________________ TestSingularities.test_singular_integral_vac_estell ___________________________________
self = <tests.test_integrals.TestSingularities object at 0x7fb16b450e90>
@pytest.mark.unit
def test_singular_integral_vac_estell(self):
"""Test calculating Bplasma for vacuum estell, which should be near 0."""
eq = get("ESTELL")
grid = LinearGrid(M=25, N=25, NFP=eq.NFP)
keys = [
"K_vc",
"B",
"|B|^2",
"R",
"phi",
"Z",
"e^rho",
"n_rho",
"|e_theta x e_zeta|",
]
data = eq.compute(keys, grid=grid)
k = min(grid.num_theta, grid.num_zeta * grid.NFP)
s = k - 1
q = k // 2 + int(np.sqrt(k))
interpolator = DFTInterpolator(grid, grid, s, q)
Bplasma = virtual_casing_biot_savart(data, data, interpolator, loop=True)
# need extra factor of B/2 bc we're evaluating on plasma surface
Bplasma += data["B"] / 2
Bplasma = np.linalg.norm(Bplasma, axis=-1)
# scale by total field magnitude
B = Bplasma / np.linalg.norm(data["B"], axis=-1).mean()
> np.testing.assert_allclose(B, 0, atol=0.005)
tests/test_integrals.py:675:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
args = (<function assert_allclose.<locals>.compare at 0x7fb16b97dda0>, array([0.3662026 , 0.25978221, 1.45680712, ..., 0.25034399, 0.25539644,
0.25843631]), array(0))
kwds = {'equal_nan': True, 'err_msg': '', 'header': 'Not equal to tolerance rtol=1e-07, atol=0.005', 'verbose': True}
@wraps(func)
def inner(*args, **kwds):
with self._recreate_cm():
> return func(*args, **kwds)
E AssertionError:
E Not equal to tolerance rtol=1e-07, atol=0.005
E
E Mismatched elements: 2601 / 2601 (100%)
E Max absolute difference: 3.92344324
E Max relative difference: inf
E x: array([0.366203, 0.259782, 1.456807, ..., 0.250344, 0.255396, 0.258436])
E y: array(0)
../../../miniconda3/envs/desc-env-pip/lib/python3.12/contextlib.py:81: AssertionError
================================================= short test summary info =================================================
FAILED tests/test_integrals.py::TestSingularities::test_singular_integral_vac_estell - AssertionError:
=========================================== 1 failed, 843 deselected in 17.47s ============================================
(desc-env-pip) kaya@fedora:~/Documents/project/DESC$
| benchmark_name | dt(%) | dt(s) | t_new(s) | t_old(s) |
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
test_build_transform_fft_lowres | -4.34 +/- 6.74 | -2.36e-02 +/- 3.66e-02 | 5.19e-01 +/- 3.2e-02 | 5.43e-01 +/- 1.7e-02 |
test_equilibrium_init_medres | -4.54 +/- 3.70 | -2.09e-01 +/- 1.71e-01 | 4.40e+00 +/- 1.4e-01 | 4.61e+00 +/- 9.5e-02 |
test_equilibrium_init_highres | -3.30 +/- 6.63 | -1.94e-01 +/- 3.90e-01 | 5.69e+00 +/- 1.1e-01 | 5.88e+00 +/- 3.7e-01 |
test_objective_compile_dshape_current | -0.50 +/- 4.49 | -2.05e-02 +/- 1.83e-01 | 4.05e+00 +/- 1.3e-01 | 4.07e+00 +/- 1.3e-01 |
test_objective_compute_dshape_current | +0.12 +/- 1.55 | +6.02e-06 +/- 7.96e-05 | 5.14e-03 +/- 6.6e-05 | 5.13e-03 +/- 4.5e-05 |
test_objective_jac_dshape_current | +1.64 +/- 6.10 | +7.19e-04 +/- 2.67e-03 | 4.45e-02 +/- 1.9e-03 | 4.38e-02 +/- 1.8e-03 |
test_perturb_2 | +6.81 +/- 6.40 | +1.34e+00 +/- 1.26e+00 | 2.10e+01 +/- 1.2e+00 | 1.97e+01 +/- 4.8e-01 |
test_proximal_freeb_jac | -3.25 +/- 1.51 | -2.41e-01 +/- 1.12e-01 | 7.17e+00 +/- 7.8e-02 | 7.41e+00 +/- 8.0e-02 |
test_solve_fixed_iter | +1.08 +/- 2.49 | +3.46e-01 +/- 7.95e-01 | 3.23e+01 +/- 2.6e-01 | 3.19e+01 +/- 7.5e-01 |
test_LinearConstraintProjection_build | -1.07 +/- 3.08 | -1.10e-01 +/- 3.16e-01 | 1.01e+01 +/- 2.6e-01 | 1.02e+01 +/- 1.8e-01 |
test_build_transform_fft_midres | +0.35 +/- 5.44 | +2.30e-03 +/- 3.57e-02 | 6.59e-01 +/- 2.3e-02 | 6.57e-01 +/- 2.7e-02 |
test_build_transform_fft_highres | -0.67 +/- 5.24 | -6.75e-03 +/- 5.30e-02 | 1.00e+00 +/- 4.4e-02 | 1.01e+00 +/- 2.9e-02 |
test_equilibrium_init_lowres | +1.23 +/- 4.73 | +5.49e-02 +/- 2.12e-01 | 4.53e+00 +/- 1.6e-01 | 4.48e+00 +/- 1.3e-01 |
test_objective_compile_atf | -1.02 +/- 1.95 | -8.94e-02 +/- 1.71e-01 | 8.71e+00 +/- 1.5e-01 | 8.80e+00 +/- 7.9e-02 |
test_objective_compute_atf | +3.97 +/- 7.08 | +6.73e-04 +/- 1.20e-03 | 1.76e-02 +/- 5.4e-04 | 1.70e-02 +/- 1.1e-03 |
test_objective_jac_atf | -0.35 +/- 2.62 | -6.94e-03 +/- 5.22e-02 | 1.98e+00 +/- 3.4e-02 | 1.99e+00 +/- 3.9e-02 |
test_perturb_1 | +0.39 +/- 1.26 | +6.03e-02 +/- 1.97e-01 | 1.56e+01 +/- 1.1e-01 | 1.56e+01 +/- 1.6e-01 |
test_proximal_jac_atf | -0.22 +/- 0.88 | -1.87e-02 +/- 7.36e-02 | 8.34e+00 +/- 5.0e-02 | 8.36e+00 +/- 5.4e-02 |
+test_proximal_freeb_compute | -7.75 +/- 2.00 | -1.62e-02 +/- 4.18e-03 | 1.92e-01 +/- 1.6e-03 | 2.09e-01 +/- 3.8e-03 |
test_solve_fixed_iter_compiled | +0.00 +/- 2.59 | +3.40e-04 +/- 5.47e-01 | 2.11e+01 +/- 3.8e-01 | 2.11e+01 +/- 3.9e-01 | |
useful for edge case of axissymetric device
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #1440 +/- ##
==========================================
- Coverage 95.69% 95.65% -0.04%
==========================================
Files 101 101
Lines 25641 25644 +3
==========================================
- Hits 24536 24529 -7
- Misses 1105 1115 +10
|
|
Actually, I cannot do that since I use the update in |
| # to the frequency content of R, Z, λ to net the best function approximation. | ||
| Nt = grid.num_theta | ||
| Nz = grid.num_zeta * grid.NFP | ||
| st = int(min(1 + jnp.sqrt(Nt), Nt)) |
There was a problem hiding this comment.
going to ask my stupid question but why this particular choice for st? it is basically always the former option (1 + sqrt(Nt)) unless Nt is 1 or 2
There was a problem hiding this comment.
It can't be empty region (>0) or more than full domain (<=Nt) and Malhotra et al. chose the sqrt
| an = -2 * jnp.flip( | ||
| take( | ||
| a.imag, | ||
| jnp.arange(1, a.shape[axis] - is_even), |
There was a problem hiding this comment.
I don't follow this logic completely, so when is_even is True, are there any sin coefficients?
There was a problem hiding this comment.
When n is even, instead of returning coefficients that can be dotted with
[sin(k𝐱), ..., sin(𝐱), 1, cos(𝐱), ..., cos(k𝐱)]
we return
[sin([k-1]𝐱), ..., sin(𝐱), 1, cos(𝐱), ..., cos(k𝐱)]
because the coefficient for sin(k𝐱) is always zero in that case. By doing this, we ensure there are always the same number of returned coefficients n as there were interpolated parameters n.
|
Another thing to check: with the changes to DFTInterpolator, is it still a lot slower? It might be competitive, especially since it allows you to use a coarser eval grid |
|
View / edit / reply to this conversation on ReviewNB unalmis commented on 2025-02-03T00:26:05Z It is weird that the optimizer does worse on the last step compared to before. |
It is competitive with courser evaluation grid. For example |
| # vmap found more efficient than fori_loop, esp on gpu, but uses more memory | ||
| if loop: | ||
| f = fori_loop(0, v.size, polar_pt_loop, f) | ||
| f = fori_loop( |
There was a problem hiding this comment.
i think also worth using vmap_chunked or batch_map or whatever here
There was a problem hiding this comment.
This is a case where chunking does not generalize the for loop because the loop reduces as the sum is computed (so max array size is 1), while chunking does not.
Should I heed your old code comment and interpret chunk size=1 as fori_loop and chunk_size!=1 as a chunked vmap? Or just the latter?
There was a problem hiding this comment.
ahh good point. Then yeah I think leave the loop for chunksize=1 but use vmap_chunked for the else
There was a problem hiding this comment.
not for this PR, but at some point might be worth implementing some sort of batched reduction desc.batching
|
recommend squash merge after someone overrides codecov |
DFTInterpolator.chunk_sizeoption to free boundary.st,szfor partition size.Bounce2Dbounce integral methods and refactors the code to reuse the Vandermonde interpolation matrix while retaining JIT's ability to fuse the matrix multiply. Benchmarking shows thiseffective rippleandGamma_c.motivation for pr (resolved by the above updates)
I have some tests to incrementally test what feature of shaping causes the issue
That should help search for the issue. For now it seems 3D shaping fails the green's id test for$\phi = 1$ . Likewise, the
FFTinterpolators andDFTinterpolators disagree in their output, indicating some bug in one or either, although the implementation here and on interpax looks correct.