Simplify normalization scales

[YigitElma]

Simplify normalizations in compute_scaling_factors to speed up the building of objectives.

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +0.10 +/- 6.96     | +5.26e-04 +/- 3.55e-02 |  5.11e-01 +/- 3.4e-02  |  5.10e-01 +/- 8.3e-03  |
 test_build_transform_fft_midres         |     -0.06 +/- 1.63     | -3.85e-04 +/- 9.68e-03 |  5.95e-01 +/- 8.7e-03  |  5.95e-01 +/- 4.3e-03  |
 test_build_transform_fft_highres        |     +1.45 +/- 2.97     | +1.43e-02 +/- 2.94e-02 |  1.00e+00 +/- 2.7e-02  |  9.89e-01 +/- 1.1e-02  |
 test_equilibrium_init_lowres            |     +0.08 +/- 3.02     | +2.91e-03 +/- 1.10e-01 |  3.65e+00 +/- 1.0e-01  |  3.65e+00 +/- 4.6e-02  |
 test_equilibrium_init_medres            |     +0.20 +/- 4.44     | +8.34e-03 +/- 1.83e-01 |  4.13e+00 +/- 1.0e-01  |  4.12e+00 +/- 1.5e-01  |
 test_equilibrium_init_highres           |     +0.79 +/- 2.40     | +4.37e-02 +/- 1.33e-01 |  5.57e+00 +/- 1.3e-01  |  5.52e+00 +/- 2.1e-02  |
 test_objective_compile_dshape_current   |     +0.42 +/- 2.11     | +1.62e-02 +/- 8.14e-02 |  3.88e+00 +/- 8.1e-02  |  3.86e+00 +/- 1.1e-02  |
 test_objective_compile_atf              |     -0.33 +/- 1.91     | -2.73e-02 +/- 1.59e-01 |  8.28e+00 +/- 8.2e-02  |  8.30e+00 +/- 1.4e-01  |
 test_objective_compute_dshape_current   |     +1.28 +/- 3.66     | +1.58e-05 +/- 4.55e-05 |  1.26e-03 +/- 3.9e-05  |  1.24e-03 +/- 2.3e-05  |
 test_objective_compute_atf              |     -0.76 +/- 4.91     | -3.24e-05 +/- 2.08e-04 |  4.20e-03 +/- 1.7e-04  |  4.24e-03 +/- 1.2e-04  |
 test_objective_jac_dshape_current       |     -0.70 +/- 8.85     | -2.70e-04 +/- 3.42e-03 |  3.84e-02 +/- 2.6e-03  |  3.86e-02 +/- 2.2e-03  |
 test_objective_jac_atf                  |     -0.21 +/- 3.20     | -4.05e-03 +/- 6.12e-02 |  1.91e+00 +/- 5.3e-02  |  1.91e+00 +/- 3.1e-02  |
+test_perturb_1                          |     -4.90 +/- 1.60     | -6.69e-01 +/- 2.18e-01 |  1.30e+01 +/- 6.8e-02  |  1.36e+01 +/- 2.1e-01  |
 test_perturb_2                          |     -4.06 +/- 1.57     | -7.57e-01 +/- 2.93e-01 |  1.79e+01 +/- 2.2e-01  |  1.87e+01 +/- 1.9e-01  |
 test_proximal_jac_atf                   |     -0.03 +/- 0.84     | -2.09e-03 +/- 6.24e-02 |  7.39e+00 +/- 2.7e-02  |  7.39e+00 +/- 5.6e-02  |
 test_proximal_freeb_compute             |     -0.92 +/- 0.91     | -1.66e-03 +/- 1.63e-03 |  1.79e-01 +/- 1.1e-03  |  1.80e-01 +/- 1.3e-03  |
 test_proximal_freeb_jac                 |     -1.17 +/- 1.88     | -8.71e-02 +/- 1.40e-01 |  7.34e+00 +/- 7.3e-02  |  7.43e+00 +/- 1.2e-01  |
 test_solve_fixed_iter                   |     -3.04 +/- 9.87     | -5.57e-01 +/- 1.81e+00 |  1.78e+01 +/- 1.6e+00  |  1.83e+01 +/- 9.1e-01  |

[YigitElma]

eq = get("ATF")

def get_lowest_mode(basis, coeffs):
    """Return lowest order coefficient (excluding m=0 modes)."""
    # lowest order modes: [0, +1, -1, +2, -2, ...]
    m_modes = np.arange(1, eq.M + 1)
    m_modes = np.vstack((m_modes, -m_modes)).flatten(order="F")
    n_modes = np.arange(eq.N + 1)
    n_modes = np.vstack((n_modes, -n_modes)).flatten(order="F")
    for n in n_modes:
        for m in m_modes:
            try:
                x = coeffs[basis.get_idx(M=m, N=n)]
                if not np.isclose(x, 0):  # mode exists and coefficient is non-zero
                    return x
            except ValueError:
                pass
    raise ValueError("No modes found, geometry is unphysical.")

def fun():
    scales = {}
    R00 = eq.Rb_lmn[eq.surface.R_basis.get_idx(M=0, N=0)]
    R10 = get_lowest_mode(eq.surface.R_basis, eq.Rb_lmn)
    Z10 = get_lowest_mode(eq.surface.Z_basis, eq.Zb_lmn)

    scales["R0"] = R00
    scales["a"] = np.sqrt(np.abs(R10 * Z10))
    scales["A"] = np.pi * scales["a"] ** 2
    scales["V"] = 2 * np.pi * scales["R0"] * scales["A"]
    scales["B_T"] = abs(eq.Psi) / scales["A"]
    if eq.iota is not None:
        iota_avg = np.mean(np.abs(eq.iota(np.linspace(0, 1, 20))))
    else:
        iota_avg = np.mean(np.abs(eq.get_profile("iota")(np.linspace(0, 1, 20))))
    if np.isclose(iota_avg, 0):
        scales["B_P"] = scales["B_T"]
    else:
        scales["B_P"] = scales["B_T"] * iota_avg
    scales["B"] = np.sqrt(scales["B_T"] ** 2 + scales["B_P"] ** 2)
    scales["I"] = scales["B_P"] * 2 * np.pi / mu_0
    scales["p"] = scales["B"] ** 2 / (2 * mu_0)
    scales["W"] = scales["p"] * scales["V"]
    scales["J"] = scales["B"] / scales["a"] / mu_0
    scales["F"] = scales["p"] / scales["a"]
    scales["f"] = scales["F"] * scales["V"]
    scales["Psi"] = abs(eq.Psi)
    scales["n"] = 1e19
    scales["T"] = scales["p"] / (scales["n"] * elementary_charge)

def fun2():
    scales = {}
    R00 = eq.Rb_lmn[eq.surface.R_basis.get_idx(M=0, N=0)]
    R10 = get_lowest_mode(eq.surface.R_basis, eq.Rb_lmn)
    Z10 = get_lowest_mode(eq.surface.Z_basis, eq.Zb_lmn)

    scales["R0"] = R00
    scales["a"] = np.sqrt(np.abs(R10 * Z10))
    scales["A"] = np.pi * scales["a"] ** 2
    scales["V"] = 2 * np.pi * scales["R0"] * scales["A"]
    scales["B_T"] = abs(eq.Psi) / scales["A"]
    iota_avg = np.mean(np.abs(eq.get_profile("iota")(np.linspace(0, 1, 20))))
    if np.isclose(iota_avg, 0):
        scales["B_P"] = scales["B_T"]
    else:
        scales["B_P"] = scales["B_T"] * iota_avg
    scales["B"] = np.sqrt(scales["B_T"] ** 2 + scales["B_P"] ** 2)
    scales["I"] = scales["B_P"] * 2 * np.pi / mu_0
    scales["p"] = scales["B"] ** 2 / (2 * mu_0)
    scales["W"] = scales["p"] * scales["V"]
    scales["J"] = scales["B"] / scales["a"] / mu_0
    scales["F"] = scales["p"] / scales["a"]
    scales["f"] = scales["F"] * scales["V"]
    scales["Psi"] = abs(eq.Psi)
    scales["n"] = 1e19
    scales["T"] = scales["p"] / (scales["n"] * elementary_charge)

def fun3():
    con = FixBoundaryR(eq)
    con.build()

%timeit -n 10 fun()   # this PR scaling factor
%timeit -n 10 fun2()  # master scaling factor (without iota check)
%timeit -n 10 fun3()  # this PR constraint build

2.27 ms ± 46 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)
176 ms ± 3.56 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
13.7 ms ± 866 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)

And constraint build (fun3()) on master takes,

190 ms ± 7.81 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

[dpanici]

Under the hood, get_profiles computes iota, so if you run your benchmark code you posted on an equilibrum that has a current profile assigned instead of an iota profile (ATF has an iota profile), master and this PR will be the same amt of time.

So I guess this PR would just speed it up for equilibria with iota profiles assigned

[dpanici]

just use current Bt scale as |B| scale

[dpanici]

I think we should make it even simpler, idt anyone will ever use B_T and B_P as the scales for a problem, so we can just do scales["B"] = abs(thing.Psi) / scales["A"] and get rid of the B_T and B_P keys entirely

[dpanici]

simplify further

[dpanici]

The test fails on CI, but seems to pass locally for me, how about you @YigitElma

[YigitElma]

The same one fails locally for me

[YigitElma]

The same test on master:

Number of parameters: 10
Number of objectives: 1
Starting optimization
Using method: scipy-trust-krylov
   Iteration     Total nfev        Cost      Cost reduction    Step norm     Optimality   
       0              1          1.334e-01                                    3.795e-03   
       1              2          1.334e-01      1.224e-05      7.180e-03      6.181e-05   
       2              3          1.334e-01      1.224e-07      1.038e-02      2.029e-06   
       3              4          1.334e-01      1.497e-11      1.680e-05      3.672e-09   
Warning: `callback` raised `StopIteration`.
         Current function value: 1.334e-01
         Total delta_x: 3.034e+00
         Iterations: 3
         Function evaluations: 4
         Gradient evaluations: 4
         Hessian evaluations: 4

On this PR:

Number of parameters: 10
Number of objectives: 1
Starting optimization
Using method: scipy-trust-krylov
   Iteration     Total nfev        Cost      Cost reduction    Step norm     Optimality   
       0              1          5.336e-01                                    1.518e-02   
       1              2          5.336e-01      4.897e-05      7.180e-03      2.472e-04   
       2              3          5.336e-01      4.895e-07      1.038e-02      8.117e-06   
       3              4          5.336e-01      5.990e-11      1.680e-05      1.469e-08   

../../miniconda3/envs/desc-env/lib/python3.12/site-packages/scipy/optimize/_trustregion.py:225: in _minimize_trust_region
    p, hits_boundary = m.solve(trust_radius)


FAILED tests/test_optimizer.py::TestAllOptimizers::test_all_optimizers_scalar[scipy-trust-krylov] - RuntimeWarning: invalid value encountered in multiply

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 95.46%. Comparing base (6ebaa23) to head (830cbb2).
Report is 1685 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #1192   +/-   ##
=======================================
  Coverage   95.45%   95.46%           
=======================================
  Files          87       87           
  Lines       22266    22268    +2     
=======================================
+ Hits        21255    21258    +3     
+ Misses       1011     1010    -1     

Files with missing lines	Coverage Δ
desc/objectives/normalization.py	96.66% <100.00%> (+0.11%)	⬆️

... and 4 files with indirect coverage changes

[YigitElma]

@f0uriest Idk why but free boundary notebook throws an error for plot_2d. I ran it locally (making sure that my matplotlib version is the same with CI) and don't see a problem. Can you check?

[ddudt]

I think some of our scales are off. Testing with approximate W7-X values of Psi = 2 and a = 0.5, this results in:

• B = 5.6 T
• p = 12.7 MPa
• T = 7.9 MeV

These values are way too high! I suggest the following changes:

• Remove the np.sqrt(2) factor from B
• Hard-code scales["T"] = 1e3 like we do for density
• Change the pressure scale to scales["p"] = 2 * elementary_charge * scales["n"] * scales["T"] (which will be 3.2 kPa)

It would also be nice if we had a way to change these hard-coded temperatures and densities. These values are fine for experiment scales but low for reactor scales. Maybe something like this could work:

scales["n"] = np.max(eq.ne_l)

(and similar for temperature). I think that would give a decent estimate independent of the Profile type

[f0uriest]

I think some of our scales are off. Testing with approximate W7-X values of Psi = 2 and a = 0.5, this results in:

* B = 5.6 T

* p = 12.7 MPa

* T = 7.9 MeV

These values are way too high! I suggest the following changes:

* Remove the `np.sqrt(2)` factor from B

* Hard-code `scales["T"] = 1e3` like we do for density

* Change the pressure scale to `scales["p"] = 2 * elementary_charge * scales["n"] * scales["T"]` (which will be 3.2 kPa)

It would also be nice if we had a way to change these hard-coded temperatures and densities. These values are fine for experiment scales but low for reactor scales. Maybe something like this could work:

scales["n"] = np.max(eq.ne_l)

(and similar for temperature). I think that would give a decent estimate independent of the Profile type

They seem ok to me. Running on the actual w7x equilibrium I get the scale for B is 3.11 T which is pretty close to the actual volume average B=2.78 T

I'd prefer to keep pressure scaling based on B, since for vacuum configs thats the only thing that makes sense.

I do like the idea of making temperature and density scales relative to the profiles if they exist (though i'd actually evaluate the profile rather than relying on the magnitude of the coefficients which may not mean the same thing for different types)

[ddudt]

They seem ok to me. Running on the actual w7x equilibrium I get the scale for B is 3.11 T which is pretty close to the actual volume average B=2.78 T

I still don't think we need the sqrt(2) factor, since I think it is trying to account for an equal contribution from the poloidal field but the majority of the field is toroidal. So this correction factor is generally too high, although we only care about order of magnitude so I'm splitting hairs here.

I'd prefer to keep pressure scaling based on B, since for vacuum configs thats the only thing that makes sense.

We need to at least add a correction factor to the existing pressure scale on the order of 1e-2, since the thermal energy is typically only a few percent of the magnetic energy.

I do like the idea of making temperature and density scales relative to the profiles if they exist (though i'd actually evaluate the profile rather than relying on the magnitude of the coefficients which may not mean the same thing for different types)

I agree that would be better, but the whole point of this PR is to avoid calling any compute functions to reduce compile times. So I was trying to think of a heuristic based only on the coefficients.

[f0uriest]

We need to at least add a correction factor to the existing pressure scale on the order of 1e-2, since the thermal energy is typically only a few percent of the magnetic energy.

That's fine, as long as we leave the scale for force as it is now, since the force error should be proportional to B^2, independent of p

I agree that would be better, but the whole point of this PR is to avoid calling any compute functions to reduce compile times. So I was trying to think of a heuristic based only on the coefficients.

Calling profile.compute is super cheap so that's probably fine (that was the initial idea of this PR anyways, and we'd only do it if the profile isn't None so in most cases its a no-op)

[f0uriest]

Update title and description before merging

[YigitElma]

Thanks for taking this over @ddudt

[YigitElma]

I cannot approve but I would 😅

I made my requested changes