Simplify normalization scales (#1192)

Simplify normalizations in `compute_scaling_factors` to speed up the building of objectives.
PlasmaControl · Aug 18, 2024 · e53ce01 · e53ce01
2 parents 6ebaa23 + 830cbb2
commit e53ce01
Showing 1 changed file with 22 additions and 14 deletions.
diff --git a/desc/objectives/normalization.py b/desc/objectives/normalization.py
@@ -38,24 +38,32 @@ def get_lowest_mode(basis, coeffs):
 
         scales["R0"] = R00
         scales["a"] = np.sqrt(np.abs(R10 * Z10))
+        scales["Psi"] = abs(thing.Psi)
         scales["A"] = np.pi * scales["a"] ** 2
         scales["V"] = 2 * np.pi * scales["R0"] * scales["A"]
-        scales["B_T"] = abs(thing.Psi) / scales["A"]
-        iota_avg = np.mean(np.abs(thing.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["B"] = scales["Psi"] / scales["A"] * 1.25
+        B_pressure = scales["B"] ** 2 / (2 * mu_0)
+        scales["I"] = scales["B"] * 2 * np.pi / mu_0
+        scales["W"] = B_pressure * scales["V"]
         scales["J"] = scales["B"] / scales["a"] / mu_0
-        scales["F"] = scales["p"] / scales["a"]
+        scales["F"] = B_pressure / scales["a"]
         scales["f"] = scales["F"] * scales["V"]
-        scales["Psi"] = abs(thing.Psi)
-        scales["n"] = 1e19
-        scales["T"] = scales["p"] / (scales["n"] * elementary_charge)
+
+        if thing.pressure is not None:
+            p0 = float(thing.pressure(0)[0])
+        else:
+            scales["n"] = float(
+                ((thing.atomic_number(0) + 1) / 2 * thing.electron_density(0))[0]
+            )
+            scales["T"] = np.mean(
+                [thing.electron_temperature(0), thing.ion_temperature(0)]
+            )
+            p0 = elementary_charge * 2 * scales["n"] * scales["T"]
+        if p0 < 1:  # vacuum
+            scales["p"] = B_pressure
+        else:
+            scales["p"] = p0
+        scales["W_p"] = scales["p"] * scales["V"] / 2
 
     elif isinstance(thing, FourierRZToroidalSurface):
         R00 = thing.R_lmn[thing.R_basis.get_idx(M=0, N=0)]