Speed up linearization construction

[nicamo]

Summary

Reduce the cost of constructing and refreshing the nonlinear-solver linearization while preserving the existing relinearization trigger and full-update policy.

This also fixes an existing Lao85 profile-derivative bug on main: when alpha_logic or beta_logic is enabled, the stored arrays already contain the dependent final coefficient. Passing those complete arrays back through initialize_profile() appends another dependent coefficient on every finite-difference profile update. The corrected path passes only the independent coefficients and lets initialize_profile() reconstruct the dependent term exactly once.

Changes

• Add n_linearization_workers to nl_solver, defaulting to 1 to preserve serial behavior.
• Build independent dIydI current columns and dIydtheta profile-parameter columns in isolated worker processes.
• Limit each construction worker to one native BLAS/OpenMP thread, making the worker count the total CPU budget for Jacobian construction.
• Restore normal native threading after construction, so ordinary timestep GS solves retain their inner parallelism.
• Reuse accepted finite-difference current perturbations on later relinearizations.
• Accept a reused perturbation after one GS solve only when its inferred amplitude is within a factor of 4/3; otherwise fall back to the original calibration path.
• Keep the first linearization and core-mask-forced linearizations on the original perturbation path.
• Start every finite-difference GS solve from the same unperturbed plasma_psi, profile state, and deterministic NK random state.
• Pass only independent Lao85 alpha/beta coefficients during profile finite differences, avoiding repeated logic-coefficient expansion already present on main.
• Document the worker argument in Examples 5a, 5b, 5c, and 11.
• Add focused tests for perturbation reuse, fallback behavior, state reset, serial/parallel dispatch, profile-parameter handling, and Lao logic coefficients.

The implementation uses multiprocessing on platforms supporting the fork start method. The serial default remains portable.

Performance

At the expensive shot-53320 step-356 checkpoint, constructing 74 current-response columns took:

• 378.32 s with one worker.
• 85.57 s with six workers.

This is a measured 4.42x speedup for that Jacobian build.

For perturbation reuse, the first full build is deliberately unchanged. In the measured replay, the clean second 74-column rebuild took 85.73 s; guarded reuse reduced it to approximately 69 s. Later gains depend on how stale the accepted perturbations have become.

A live shot-53320 replay through step 399 and five relinearizations remained within:

• 3.08 mm maximum plasma-descriptor deviation.
• 91.84 A maximum physical active-coil deviation.

These figures isolate the linearization improvements. End-to-end gains depend on how much of a run is spent rebuilding Jacobians.

Compatibility

• n_linearization_workers=1 preserves the existing serial route.
• The relinearization trigger and all-column update policy are unchanged.
• The first linearization uses the original perturbation estimates.
• Rejected reused perturbations automatically fall back to the original two-stage calibration.
• Normal timestep GS threading is not restricted by the construction worker limit.
• Lao85 profiles with alpha_logic and/or beta_logic now retain only their intended independent coefficients during finite-difference construction.

Validation

• python -m black --check freegsnke/
• python -m isort --check --diff freegsnke/
• python -m pytest -q freegsnke/tests/test_linearisation_perturbations.py

Result: 13 passed.

[kpentland]

Multithreading was so needed here! Makes things much faster!