1use crate::vacuum::calculate_vacuum_field;
13use fusion_math::interp::{gradient_2d, interp2d};
14use fusion_math::linalg::eig_2x2;
15use fusion_types::config::ReactorConfig;
16use fusion_types::constants::MU0_SI;
17use fusion_types::error::{FusionError, FusionResult};
18use fusion_types::state::{Grid2D, StabilityResult};
19use ndarray::Array2;
20
21const MINOR_RADIUS_RATIO: f64 = 3.0;
23
24const INTERNAL_INDUCTANCE: f64 = 0.8;
26
27const BETA_POLOIDAL: f64 = 0.5;
29
30const FORCE_PERTURBATION: f64 = 0.01;
32
33const MAX_STABLE_DECAY_INDEX: f64 = 1.5;
35
36const FORCE_TOLERANCE: f64 = 1e4;
38
39const MAX_NEWTON_ITER: usize = 10;
41
42const COIL_PERTURBATION_MA: f64 = 0.1;
44
45const MAX_CURRENT_CORRECTION_MA: f64 = 5.0;
47
48const MIN_RADIUS_M: f64 = 1e-6;
49
50fn validate_stability_inputs(psi: &Array2<f64>, grid: &Grid2D, r: f64, z: f64) -> FusionResult<()> {
51 if psi.nrows() != grid.nz || psi.ncols() != grid.nr {
52 return Err(FusionError::ConfigError(format!(
53 "stability psi shape mismatch: expected ({}, {}), got ({}, {})",
54 grid.nz,
55 grid.nr,
56 psi.nrows(),
57 psi.ncols()
58 )));
59 }
60 if psi.iter().any(|v| !v.is_finite()) {
61 return Err(FusionError::ConfigError(
62 "stability psi must contain only finite values".to_string(),
63 ));
64 }
65 if !r.is_finite() || r <= MIN_RADIUS_M {
66 return Err(FusionError::ConfigError(format!(
67 "stability radius r must be finite and > {MIN_RADIUS_M}, got {r}"
68 )));
69 }
70 if !z.is_finite() {
71 return Err(FusionError::ConfigError(format!(
72 "stability vertical coordinate z must be finite, got {z}"
73 )));
74 }
75 Ok(())
76}
77
78pub fn decay_index(psi: &Array2<f64>, grid: &Grid2D, r: f64, z: f64) -> FusionResult<f64> {
84 validate_stability_inputs(psi, grid, r, z)?;
85
86 let (_dpsi_dz, dpsi_dr) = gradient_2d(psi, grid);
87
88 let bz = interp2d(&dpsi_dr, grid, r, z) / r;
90 if !bz.is_finite() {
91 return Err(FusionError::ConfigError(
92 "decay_index computed non-finite B_Z".to_string(),
93 ));
94 }
95
96 if bz.abs() < 1e-12 {
97 return Err(FusionError::ConfigError(
98 "decay_index undefined when |B_Z| is near zero".to_string(),
99 ));
100 }
101
102 let eps = FORCE_PERTURBATION;
104 if r <= eps {
105 return Err(FusionError::ConfigError(format!(
106 "decay_index requires r > perturbation epsilon ({eps})"
107 )));
108 }
109 let bz_plus = interp2d(&dpsi_dr, grid, r + eps, z) / (r + eps);
110 let bz_minus = interp2d(&dpsi_dr, grid, r - eps, z) / (r - eps);
111 if !bz_plus.is_finite() || !bz_minus.is_finite() {
112 return Err(FusionError::ConfigError(
113 "decay_index finite-difference B_Z samples became non-finite".to_string(),
114 ));
115 }
116 let dbz_dr = (bz_plus - bz_minus) / (2.0 * eps);
117 if !dbz_dr.is_finite() {
118 return Err(FusionError::ConfigError(
119 "decay_index dB_Z/dR became non-finite".to_string(),
120 ));
121 }
122
123 let n = -(r / bz) * dbz_dr;
124 if !n.is_finite() {
125 return Err(FusionError::ConfigError(
126 "decay_index result became non-finite".to_string(),
127 ));
128 }
129 Ok(n)
130}
131
132pub fn calculate_forces(
140 psi: &Array2<f64>,
141 grid: &Grid2D,
142 r: f64,
143 z: f64,
144 i_plasma_ma: f64,
145) -> FusionResult<(f64, f64)> {
146 validate_stability_inputs(psi, grid, r, z)?;
147 if !i_plasma_ma.is_finite() {
148 return Err(FusionError::ConfigError(format!(
149 "i_plasma_ma must be finite, got {i_plasma_ma}"
150 )));
151 }
152
153 let i_amp = i_plasma_ma * 1e6;
154 if !i_amp.is_finite() {
155 return Err(FusionError::ConfigError(
156 "plasma current conversion to amperes became non-finite".to_string(),
157 ));
158 }
159 let (dpsi_dz, dpsi_dr) = gradient_2d(psi, grid);
160
161 let bz = interp2d(&dpsi_dr, grid, r, z) / r;
163 let br = -interp2d(&dpsi_dz, grid, r, z) / r;
164 if !bz.is_finite() || !br.is_finite() {
165 return Err(FusionError::ConfigError(
166 "calculate_forces computed non-finite B-field components".to_string(),
167 ));
168 }
169
170 let a = r / MINOR_RADIUS_RATIO;
172 let shafranov_term = (8.0 * r / a).ln() + BETA_POLOIDAL + INTERNAL_INDUCTANCE / 2.0 - 1.5;
173 let f_hoop = (MU0_SI * i_amp * i_amp / 2.0) * shafranov_term / r;
174
175 let f_lorentz_r = i_amp * bz * 2.0 * std::f64::consts::PI * r;
177 let f_lorentz_z = -i_amp * br * 2.0 * std::f64::consts::PI * r;
178
179 let f_radial = f_hoop + f_lorentz_r;
180 let f_vertical = f_lorentz_z;
181 if !f_radial.is_finite() || !f_vertical.is_finite() {
182 return Err(FusionError::ConfigError(
183 "calculate_forces produced non-finite force values".to_string(),
184 ));
185 }
186
187 Ok((f_radial, f_vertical))
188}
189
190pub fn analyze_stability(
195 psi: &Array2<f64>,
196 grid: &Grid2D,
197 r_eq: f64,
198 z_eq: f64,
199 i_plasma_ma: f64,
200) -> FusionResult<StabilityResult> {
201 validate_stability_inputs(psi, grid, r_eq, z_eq)?;
202 if !i_plasma_ma.is_finite() {
203 return Err(FusionError::ConfigError(format!(
204 "i_plasma_ma must be finite, got {i_plasma_ma}"
205 )));
206 }
207 let eps = FORCE_PERTURBATION;
208
209 let (fr0, fz0) = calculate_forces(psi, grid, r_eq, z_eq, i_plasma_ma)?;
211
212 let (fr_rp, fz_rp) = calculate_forces(psi, grid, r_eq + eps, z_eq, i_plasma_ma)?;
214 let (fr_rm, fz_rm) = calculate_forces(psi, grid, r_eq - eps, z_eq, i_plasma_ma)?;
215
216 let (fr_zp, fz_zp) = calculate_forces(psi, grid, r_eq, z_eq + eps, i_plasma_ma)?;
218 let (fr_zm, fz_zm) = calculate_forces(psi, grid, r_eq, z_eq - eps, i_plasma_ma)?;
219
220 let k_rr = -(fr_rp - fr_rm) / (2.0 * eps);
222 let k_rz = -(fr_zp - fr_zm) / (2.0 * eps);
223 let k_zr = -(fz_rp - fz_rm) / (2.0 * eps);
224 let k_zz = -(fz_zp - fz_zm) / (2.0 * eps);
225
226 let k_matrix = [[k_rr, k_rz], [k_zr, k_zz]];
227 let (eigenvalues, eigenvectors) = eig_2x2(&k_matrix);
228
229 let n = decay_index(psi, grid, r_eq, z_eq)?;
230 let is_stable = n > 0.0 && n < MAX_STABLE_DECAY_INDEX;
231
232 let result = StabilityResult {
233 eigenvalues,
234 eigenvectors,
235 decay_index: n,
236 radial_force_mn: fr0 / 1e6,
237 vertical_force_mn: fz0 / 1e6,
238 is_stable,
239 };
240 if !result.eigenvalues[0].is_finite()
241 || !result.eigenvalues[1].is_finite()
242 || !result.decay_index.is_finite()
243 || !result.radial_force_mn.is_finite()
244 || !result.vertical_force_mn.is_finite()
245 {
246 return Err(FusionError::ConfigError(
247 "analyze_stability produced non-finite outputs".to_string(),
248 ));
249 }
250
251 Ok(result)
252}
253
254pub fn solve_force_balance(
261 config: &mut ReactorConfig,
262 grid: &Grid2D,
263 r_target: f64,
264 z_target: f64,
265 control_coil_indices: &[usize],
266) -> FusionResult<(usize, f64)> {
267 if !r_target.is_finite() || !z_target.is_finite() {
268 return Err(FusionError::ConfigError(
269 "force-balance target coordinates must be finite".to_string(),
270 ));
271 }
272 if control_coil_indices.is_empty() {
273 return Err(FusionError::ConfigError(
274 "force-balance requires at least one control coil index".to_string(),
275 ));
276 }
277
278 let mu0 = config.physics.vacuum_permeability;
279 let i_plasma_ma = config.physics.plasma_current_target / 1e6;
280 if !mu0.is_finite() || !i_plasma_ma.is_finite() {
281 return Err(FusionError::ConfigError(
282 "force-balance physics constants must be finite".to_string(),
283 ));
284 }
285
286 for iter in 0..MAX_NEWTON_ITER {
287 let psi = calculate_vacuum_field(grid, &config.coils, mu0)?;
289 let (f_r, _f_z) = calculate_forces(&psi, grid, r_target, z_target, i_plasma_ma)?;
290
291 if f_r.abs() < FORCE_TOLERANCE {
292 return Ok((iter, f_r.abs()));
293 }
294
295 let n_controls = control_coil_indices.len();
297 let mut jacobian = vec![0.0; n_controls];
298
299 for (j, &coil_idx) in control_coil_indices.iter().enumerate() {
300 if coil_idx < config.coils.len() {
301 let original = config.coils[coil_idx].current;
302 config.coils[coil_idx].current = original + COIL_PERTURBATION_MA;
303
304 let psi_pert = calculate_vacuum_field(grid, &config.coils, mu0)?;
305 let (f_r_pert, _) =
306 calculate_forces(&psi_pert, grid, r_target, z_target, i_plasma_ma)?;
307
308 jacobian[j] = (f_r_pert - f_r) / COIL_PERTURBATION_MA;
309 config.coils[coil_idx].current = original;
310 }
311 }
312
313 let j_total: f64 = jacobian.iter().sum();
315 if j_total.abs() < 1e-10 {
316 break;
317 }
318
319 let delta_i = (-f_r / j_total).clamp(-MAX_CURRENT_CORRECTION_MA, MAX_CURRENT_CORRECTION_MA);
320
321 for &coil_idx in control_coil_indices {
322 if coil_idx < config.coils.len() {
323 config.coils[coil_idx].current += delta_i;
324 }
325 }
326 }
327
328 let psi = calculate_vacuum_field(grid, &config.coils, mu0)?;
330 let (f_r, _) = calculate_forces(&psi, grid, r_target, z_target, i_plasma_ma)?;
331 Ok((MAX_NEWTON_ITER, f_r.abs()))
332}
333
334#[cfg(test)]
335mod tests {
336 use super::*;
337 use fusion_types::config::ReactorConfig;
338 use std::path::PathBuf;
339
340 fn project_root() -> PathBuf {
341 PathBuf::from(env!("CARGO_MANIFEST_DIR"))
342 .join("..")
343 .join("..")
344 .join("..")
345 }
346
347 fn config_path(relative: &str) -> String {
348 project_root().join(relative).to_string_lossy().to_string()
349 }
350
351 #[test]
352 fn test_decay_index_finite() {
353 let cfg = ReactorConfig::from_file(&config_path("iter_config.json")).unwrap();
354 let grid = cfg.create_grid();
355 let psi = calculate_vacuum_field(&grid, &cfg.coils, cfg.physics.vacuum_permeability)
356 .expect("valid vacuum-field inputs");
357
358 let n = decay_index(&psi, &grid, 6.2, 0.0).expect("valid finite decay-index inputs");
359 assert!(n.is_finite(), "Decay index should be finite: {n}");
360 }
361
362 #[test]
363 fn test_forces_finite() {
364 let cfg = ReactorConfig::from_file(&config_path("iter_config.json")).unwrap();
365 let grid = cfg.create_grid();
366 let psi = calculate_vacuum_field(&grid, &cfg.coils, cfg.physics.vacuum_permeability)
367 .expect("valid vacuum-field inputs");
368
369 let (fr, fz) =
370 calculate_forces(&psi, &grid, 6.2, 0.0, 15.0).expect("valid finite force inputs");
371 assert!(fr.is_finite(), "Radial force should be finite: {fr}");
372 assert!(fz.is_finite(), "Vertical force should be finite: {fz}");
373 }
374
375 #[test]
376 fn test_stability_analysis() {
377 let cfg = ReactorConfig::from_file(&config_path("iter_config.json")).unwrap();
378 let grid = cfg.create_grid();
379 let psi = calculate_vacuum_field(&grid, &cfg.coils, cfg.physics.vacuum_permeability)
380 .expect("valid vacuum-field inputs");
381
382 let result = analyze_stability(&psi, &grid, 6.2, 0.0, 15.0)
383 .expect("valid finite stability-analysis inputs");
384 assert!(
385 result.eigenvalues[0].is_finite() && result.eigenvalues[1].is_finite(),
386 "Eigenvalues should be finite"
387 );
388 assert!(
389 result.decay_index.is_finite(),
390 "Decay index should be finite"
391 );
392 }
393
394 #[test]
395 fn test_force_balance_runs() {
396 let mut cfg =
397 ReactorConfig::from_file(&config_path("validation/iter_validated_config.json"))
398 .unwrap();
399 let grid = cfg.create_grid();
400
401 let (iters, final_force) = solve_force_balance(&mut cfg, &grid, 6.2, 0.0, &[2, 3])
403 .expect("valid finite force-balance inputs");
404
405 assert!(iters <= 10, "Should not exceed max iterations");
407 assert!(final_force.is_finite(), "Final force should be finite");
408 }
409
410 #[test]
411 fn test_stability_rejects_invalid_runtime_inputs() {
412 let cfg = ReactorConfig::from_file(&config_path("iter_config.json")).unwrap();
413 let grid = cfg.create_grid();
414 let psi = calculate_vacuum_field(&grid, &cfg.coils, cfg.physics.vacuum_permeability)
415 .expect("valid vacuum-field inputs");
416
417 assert!(decay_index(&psi, &grid, f64::NAN, 0.0).is_err());
418 assert!(calculate_forces(&psi, &grid, 0.0, 0.0, 15.0).is_err());
419 assert!(analyze_stability(&psi, &grid, 0.0, 0.0, 15.0).is_err());
420
421 let mut cfg2 = ReactorConfig::from_file(&config_path("iter_config.json")).unwrap();
422 assert!(solve_force_balance(&mut cfg2, &grid, 6.2, 0.0, &[]).is_err());
423 }
424}