feat(vmec): add find_theta inverse method to VMECGeometry

[krystophny]

User description

Summary

• Add Newton-Raphson inversion to find VMEC poloidal angle from (R, Z, zeta) at a given flux surface
• Uses analytic Jacobian from Fourier derivatives for efficiency
• Supports single guess with fast Newton iteration and multi-start fallback for robust global convergence
• Vectorized batch processing for arrays of points

Use Case

Computing VMEC theta coordinates for particle loss footprints from BEAMS3D simulations where only (R, Z, phi) are available at the wall crossing.

Test plan

• Roundtrip test: forward (theta → R, Z) then inverse (R, Z → theta) recovers original theta within 1e-10 tolerance

---

PR Type

Enhancement

---

Description

• Add find_theta method to invert VMEC poloidal angle from (R, Z) coordinates

• Implements Newton-Raphson iteration with analytic Jacobian from Fourier derivatives

• Supports single initial guess for fast convergence or multi-start fallback

• Handles vectorized batch processing for arrays of points

---

Diagram Walkthrough

flowchart LR
  A["(R, Z, zeta) targets"] --> B["Interpolate coefficients at surface s"]
  B --> C["Single guess provided?"]
  C -->|Yes| D["Newton-Raphson from guess"]
  C -->|No| E["Multi-start with n_starts guesses"]
  D --> F["Analytic Jacobian evaluation"]
  E --> F
  F --> G["Gauss-Newton step minimization"]
  G --> H["theta in [0, 2π)"]

Loading

File Walkthrough

Relevant files

Enhancement

vmec.py Implement VMEC theta inversion with Newton-Raphson              python/libneo/vmec.py Add find_theta method to VMECGeometry class for inverting poloidal angle from (R, Z) coordinates Implements Newton-Raphson iteration with analytic Jacobian computed from Fourier derivatives Supports both single initial guess mode and multi-start fallback for robust global convergence Includes vectorized batch processing for arrays of target points with linear interpolation of Fourier coefficients at arbitrary flux surfaces +135/-0

---

[qodo-code-review]

PR Compliance Guide 🔍

Below is a summary of compliance checks for this PR:

Security Compliance
🟢	No security concerns identified No security vulnerabilities detected by AI analysis. Human verification advised for critical code.
Ticket Compliance
⚪	🎫 No ticket provided Create ticket/issue
Codebase Duplication Compliance
⚪	Codebase context is not defined Follow the guide to enable codebase context checks.
Custom Compliance
🟢	Generic: Comprehensive Audit Trails Objective: To create a detailed and reliable record of critical system actions for security analysis and compliance. Status: Passed Learn more about managing compliance generic rules or creating your own custom rules
	Generic: Meaningful Naming and Self-Documenting Code Objective: Ensure all identifiers clearly express their purpose and intent, making code self-documenting Status: Passed Learn more about managing compliance generic rules or creating your own custom rules
	Generic: Secure Error Handling Objective: To prevent the leakage of sensitive system information through error messages while providing sufficient detail for internal debugging. Status: Passed Learn more about managing compliance generic rules or creating your own custom rules
	Generic: Secure Logging Practices Objective: To ensure logs are useful for debugging and auditing without exposing sensitive information like PII, PHI, or cardholder data. Status: Passed Learn more about managing compliance generic rules or creating your own custom rules
🔴	Generic: Robust Error Handling and Edge Case Management Objective: Ensure comprehensive error handling that provides meaningful context and graceful degradation Status: Missing parameter validation: find_theta does not validate key boundary/edge inputs (e.g., n_starts can be 0 leading to a silent default result, and s is not range-checked), which can cause incorrect outputs without actionable errors. Referred Code def find_theta( self, R_target: float | np.ndarray, Z_target: float | np.ndarray, zeta: float | np.ndarray, s: float = 1.0, *, theta_guess: float | np.ndarray | None = None, max_iter: int = 20, tol: float = 1.0e-10, n_starts: int = 12, use_asym: bool = True, ) -> np.ndarray: """ Find VMEC poloidal angle theta for given (R, Z) at fixed (s, zeta). Uses Newton-Raphson iteration with analytic Jacobian from Fourier derivatives. Falls back to multi-start if no initial guess is provided. Parameters ---------- ... (clipped 112 lines) Learn more about managing compliance generic rules or creating your own custom rules
	Generic: Security-First Input Validation and Data Handling Objective: Ensure all data inputs are validated, sanitized, and handled securely to prevent vulnerabilities Status: Unvalidated external inputs: find_theta accepts externally-provided numerical inputs (notably s, max_iter, tol, and n_starts) without validating safe/expected ranges, enabling erroneous behavior (e.g., n_starts <= 0) rather than failing safely with clear context. Referred Code def find_theta( self, R_target: float | np.ndarray, Z_target: float | np.ndarray, zeta: float | np.ndarray, s: float = 1.0, *, theta_guess: float | np.ndarray | None = None, max_iter: int = 20, tol: float = 1.0e-10, n_starts: int = 12, use_asym: bool = True, ) -> np.ndarray: """ Find VMEC poloidal angle theta for given (R, Z) at fixed (s, zeta). Uses Newton-Raphson iteration with analytic Jacobian from Fourier derivatives. Falls back to multi-start if no initial guess is provided. Parameters ---------- ... (clipped 113 lines) Learn more about managing compliance generic rules or creating your own custom rules

Compliance status legend

🟢 - Fully Compliant
🟡 - Partial Compliant
🔴 - Not Compliant
⚪ - Requires Further Human Verification
🏷️ - Compliance label

[qodo-code-review]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Impact
High-level	Vectorize the implementation for performance The current implementation iterates over input points using a slow Python loop, despite the PR description claiming vectorized processing. Refactor the solver to use vectorized NumPy operations for significant performance gains on large arrays. Examples: python/libneo/vmec.py [497-513] for i in range(n): R_t, Z_t, z = float(R_arr[i]), float(Z_arr[i]), float(zeta_arr[i]) if guess_arr is not None: # Single Newton from provided guess th, _ = _newton(float(guess_arr[i]), z, R_t, Z_t) theta_out[i] = th else: # Multi-start: try n_starts initial guesses best_th = 0.0 best_err = float("inf") ... (clipped 7 lines) Solution Walkthrough: Before: def find_theta(self, R_target, Z_target, zeta, ...): R_arr = np.atleast_1d(R_target) Z_arr = np.atleast_1d(Z_target) zeta_arr = np.atleast_1d(zeta) n = R_arr.size theta_out = np.zeros(n) def _newton(th0, z, R_t, Z_t): # ... scalar Newton iteration ... return th # Process points one by one in a Python loop for i in range(n): R_t, Z_t, z = R_arr[i], Z_arr[i], zeta_arr[i] if guess_is_provided: theta_out[i] = _newton(guess[i], z, R_t, Z_t) else: # multi-start loop for a single point ... theta_out[i] = best_th return theta_out After: def find_theta(self, R_target, Z_target, zeta, ...): R_t = np.asarray(R_target) Z_t = np.asarray(Z_target) z = np.asarray(zeta) # Handle single vs. multi-start guesses for all points th = setup_initial_guesses_vectorized(R_t.shape) # Newton iteration loop is now the outer loop for _ in range(max_iter): # All calculations are vectorized over all points R, Z, dR, dZ = _eval_rz_vectorized(th, z) err_R = R - R_t err_Z = Z - Z_t # Vectorized Gauss-Newton step grad = 2 * (err_R * dR + err_Z * dZ) hess = 2 * (dR * dR + dZ * dZ) th -= grad / hess # Check for convergence across all points ... # Select best result for each point if using multi-start return final_thetas Suggestion importance[1-10]: 9 __ Why: The suggestion correctly identifies a major performance bottleneck due to a Python loop over input points, which contradicts the PR's stated goal of "vectorized batch processing" and severely limits the method's utility for large datasets.	High
Possible issue	Respect use_asym flag for ns==1 In the ns == 1 case, respect the use_asym flag by conditionally loading asymmetric coefficients, ensuring consistent behavior with the ns > 1 case. python/libneo/vmec.py [431-435] if ns == 1: rmnc_s = self.rmnc[:, 0] zmns_s = self.zmns[:, 0] - rmns_s = self.rmns[:, 0] if self.rmns is not None else None - zmnc_s = self.zmnc[:, 0] if self.zmnc is not None else None + if use_asym and self.rmns is not None and self.zmnc is not None: + rmns_s = self.rmns[:, 0] + zmnc_s = self.zmnc[:, 0] + else: + rmns_s = None + zmnc_s = None Apply / Chat Suggestion importance[1-10]: 8 __ Why: This suggestion fixes a bug where the use_asym flag was ignored for the ns == 1 case, leading to incorrect calculations if asymmetric terms were present but not requested.	Medium
	Validate multi-start parameter is positive Add a ValueError check at the beginning of the function to ensure n_starts is a positive integer when theta_guess is not provided, preventing silent failures. python/libneo/vmec.py [497-513] ++ if guess_arr is None and n_starts <= 0: ++ raise ValueError("n_starts must be a positive integer for multi-start mode") ++ + for i in range(n): + R_t, Z_t, z = float(R_arr[i]), float(Z_arr[i]), float(zeta_arr[i]) + if guess_arr is not None: + # Single Newton from provided guess + th, _ = _newton(float(guess_arr[i]), z, R_t, Z_t) + theta_out[i] = th + else: + # Multi-start: try n_starts initial guesses + best_th = 0.0 + best_err = float("inf") + for k in range(n_starts): + th0 = TWOPI * k / n_starts + th, err2 = _newton(th0, z, R_t, Z_t) + if err2 < best_err: + best_err = err2 + best_th = th + theta_out[i] = best_th [To ensure code accuracy, apply this suggestion manually] Suggestion importance[1-10]: 7 __ Why: This suggestion correctly identifies an edge case where a non-positive n_starts would lead to silently returning an incorrect result, and proposes a robust fix by adding input validation.	Medium
General	Return scalar for scalar input Modify the function to return a scalar float instead of a single-element array when the inputs were scalars, improving the user interface. python/libneo/vmec.py [515] +if theta_out.size == 1: + return float(theta_out[0]) return theta_out Apply / Chat Suggestion importance[1-10]: 6 __ Why: The suggestion improves the function's usability by making the return type match the input type (scalar in, scalar out), which is a common and convenient design pattern.	Low
	Simplify Gauss-Newton step calculation In the _newton function, remove the redundant 2.0 factor from both the grad and hess calculations as it cancels out. python/libneo/vmec.py [467-485] + def _newton(th0: float, z: float, R_t: float, Z_t: float) -> Tuple[float, float]: + """Newton iteration from initial guess, returns (theta, error).""" + th = th0 + for _ in range(max_iter): + R, Z, dR, dZ = _eval_rz(th, z) + err_R = R - R_t + err_Z = Z - Z_t + err2 = err_R * err_R + err_Z * err_Z + if err2 < tol: + return th % TWOPI, err2 + # Gauss-Newton step: minimize ||f||^2 -+ grad = 2.0 * (err_R * dR + err_Z * dZ) -+ hess = 2.0 * (dR * dR + dZ * dZ) ++ grad = err_R * dR + err_Z * dZ ++ hess = dR * dR + dZ * dZ + if abs(hess) < 1.0e-14: + break + th -= grad / hess + R, Z, _, _ = _eval_rz(th, z) + err2 = (R - R_t) ** 2 + (Z - Z_t) ** 2 + return th % TWOPI, err2 [To ensure code accuracy, apply this suggestion manually] Suggestion importance[1-10]: 3 __ Why: The suggestion correctly identifies that the 2.0 factor in the grad and hess calculations is redundant and can be removed, which is a minor code simplification and micro-optimization.	Low
More