Source code for scpn_fusion.core.force_balance

# SPDX-License-Identifier: AGPL-3.0-or-later | Commercial license available
# © Concepts 1996–2026 Miroslav Šotek. All rights reserved.
# © Code 2020–2026 Miroslav Šotek. All rights reserved.
# ORCID: 0009-0009-3560-0851
# Contact: www.anulum.li | protoscience@anulum.li
# SCPN Fusion Core — Force Balance
from __future__ import annotations

import numpy as np
import json
from pathlib import Path

try:
    from scpn_fusion.core._rust_compat import FusionKernel
except ImportError:
    pass
from scpn_fusion.core.stability_analyzer import StabilityAnalyzer




class ForceBalanceSolver:
    """
    Automatic Newton-Raphson Solver to find Coil Currents that result in
    ZERO Radial Force on the plasma at the target Radius.
    This guarantees static equilibrium at R_target.
    """

    def __init__(self, config_path):
        self.config_path = config_path
        self.analyzer = StabilityAnalyzer(config_path)



    def solve_for_equilibrium(
        self,
        target_R=6.2,
        target_Z=0.0,
        *,
        max_iterations: int = 10,
        jacobian_floor: float = 1e-12,
    ):
        print("--- FORCE BALANCE SOLVER (Newton-Raphson) ---")
        print(f"Target Equilibrium: R={target_R}m, Z={target_Z}m")

        # We control PF3 and PF4 (Outer coils) to control Radial Position
        # We assume Up-Down symmetry for now (PF3=PF4)
        control_indices = [2, 3]  # PF3, PF4 indexes in coil list

        # Load Physics Params
        Ip = self.analyzer.kernel.cfg["physics"]["plasma_current_target"]

        # Newton Loop
        for iter in range(max_iterations):
            # 1. Calculate Current Force
            Fr, Fz, n_idx = self.analyzer.calculate_forces(target_R, target_Z, Ip)
            print(f"Iter {iter}: Radial Force = {Fr / 1e6:.2f} MN")

            if abs(Fr) < 1e4:  # Tolerance 10 kN
                print("  -> CONVERGED. Force Balance Achieved.")
                break

            # 2. Calculate Jacobian dF/dI (Numerical Derivative)
            # We perturb PF3/PF4 together
            dI = 0.1  # MA perturbation

            # Apply perturbation
            currents = [c["current"] for c in self.analyzer.kernel.cfg["coils"]]
            original_I3 = currents[2]

            # Modify config in memory
            self.analyzer.kernel.cfg["coils"][2]["current"] += dI
            self.analyzer.kernel.cfg["coils"][3]["current"] += dI

            # Recalculate Vacuum Field (Expensive part, but necessary)
            self.analyzer.Psi_vac = self.analyzer.kernel.calculate_vacuum_field()

            # Calculate new Force
            Fr_new, _, _ = self.analyzer.calculate_forces(target_R, target_Z, Ip)

            # Jacobian J = dF/dI (MN / MA)
            J = (Fr_new - Fr) / dI
            print(f"  Jacobian dF/dI_PF34: {J / 1e6:.2f} MN/MA")

            # 3. Newton Step: I_new = I_old - F / J
            # We want F_target = 0
            # 0 = F_old + J * delta_I
            # delta_I = - F_old / J
            if not np.isfinite(J) or abs(J) < jacobian_floor:
                # Singular/ill-conditioned update; take a conservative directional step.
                delta_I = -0.25 * np.sign(Fr) if Fr != 0.0 else 0.0
                print(
                    "  Jacobian near-singular; applying conservative "
                    f"directional step ΔI={delta_I:.3f} MA"
                )
            else:
                delta_I = -Fr / J

            # Safety Clamp (don't jump more than 5 MA at once)
            delta_I = np.clip(delta_I, -5.0, 5.0)

            print(f"  Correction: Delta I = {delta_I:.3f} MA")

            # Apply correction
            new_I = original_I3 + delta_I

            # Reset Config to new values
            self.analyzer.kernel.cfg["coils"][2]["current"] = new_I
            self.analyzer.kernel.cfg["coils"][3]["current"] = new_I

            # Update Vacuum field for next iteration base
            self.analyzer.Psi_vac = self.analyzer.kernel.calculate_vacuum_field()

        # Save Result
        self.save_config()




    def save_config(self, output_path: str | Path | None = None):
        # Save to validation folder by default.
        repo_root = Path(__file__).resolve().parents[3]
        out_path = (
            Path(output_path)
            if output_path is not None
            else (repo_root / "validation" / "iter_force_balanced.json")
        )
        out_path.parent.mkdir(parents=True, exist_ok=True)
        out_path.write_text(
            json.dumps(self.analyzer.kernel.cfg, indent=4),
            encoding="utf-8",
        )
        print(f"Balanced Config Saved: {out_path}")




if __name__ == "__main__":
    # Start from the last "best" config (Genetic or Validated)
    base_cfg = "03_CODE/SCPN-Fusion-Core/validation/iter_validated_config.json"

    solver = ForceBalanceSolver(base_cfg)
    solver.solve_for_equilibrium(target_R=6.2)