I/O Utilities

Serialization of NLSQ and MCMC results to JSON, NPZ, and diagnostic output formats. The heterodyne.optimization.cmc.io module provides additional shard-level CMC I/O; see CMC (Bayesian).

JSON Utilities

JSON-safe serialization helpers for JAX arrays, NumPy scalars, and custom result types.

JSON serialization utilities for JAX arrays and numpy types.

heterodyne.io.json_utils.json_safe(obj)[source]

Convert object to JSON-serializable form.

Handles: - JAX arrays -> lists - numpy arrays -> lists - numpy scalars -> Python scalars - complex numbers -> {“real”: …, “imag”: …} dicts - Path objects -> strings - datetime -> ISO format strings - Nested dicts/lists recursively

Parameters:

obj (Any) – Object to convert

Return type:

Any

Returns:

JSON-serializable equivalent

heterodyne.io.json_utils.json_serializer(obj)[source]

Serialize object to JSON string with pretty formatting.

Parameters:

obj (Any) – Object to serialize

Return type:

str

Returns:

Pretty-printed JSON string

heterodyne.io.json_utils.load_json(path)[source]

Load JSON file.

Parameters:

path (Path | str) – Path to JSON file

Return type:

dict[str, Any]

Returns:

Parsed JSON data

heterodyne.io.json_utils.save_json(data, path)[source]

Save data to JSON file with pretty formatting.

Uses atomic write (write-to-temp + rename) to prevent partial writes.

Parameters:

  • data (Any) – Data to save

  • path (Path | str) – Output path

Return type:

None

NLSQ Writers

Writers for NLSQ optimization results.

heterodyne.io.nlsq_writers.save_nlsq_json_files(result, output_dir, prefix='nlsq')[source]

Save NLSQ results to JSON files.

Creates: - {prefix}_parameters.json: Fitted parameter values and uncertainties - {prefix}_metadata.json: Fit statistics and convergence info

Parameters:

  • result (NLSQResult) – NLSQ fitting result

  • output_dir (Path | str) – Output directory

  • prefix (str) – Filename prefix

Return type:

dict[str, Path]

Returns:

Dict mapping file type to saved path

heterodyne.io.nlsq_writers.save_nlsq_npz_file(result, output_path, include_residuals=True, include_jacobian=False, c2_exp=None)[source]

Save NLSQ results to compressed NPZ file.

NPZ format is efficient for large arrays (correlation matrices, residuals).

Parameters:

  • result (NLSQResult) – NLSQ fitting result

  • output_path (Path | str) – Output file path

  • include_residuals (bool) – Whether to include residual array

  • include_jacobian (bool) – Whether to include Jacobian matrix (large)

  • c2_exp (ndarray | None) – Experimental correlation data used to compute residuals_normalized = residuals / (0.05 * c2_exp). Saved alongside raw residuals when provided.

Return type:

Path

Returns:

Path to saved file

heterodyne.io.nlsq_writers.load_nlsq_npz_file(path)[source]

Load NLSQResult from an NPZ file saved by save_nlsq_npz_file.

Parameters:

path (Path | str) – Path to the .npz file.

Return type:

NLSQResult

Returns:

Reconstructed NLSQResult with available fields.

Raises:

FileNotFoundError – If path does not exist.

heterodyne.io.nlsq_writers.format_nlsq_summary(result)[source]

Format NLSQ result as human-readable summary.

Parameters:

result (NLSQResult) – NLSQ fitting result

Return type:

str

Returns:

Formatted summary string

MCMC Writers

Writers for MCMC/CMC analysis results.

heterodyne.io.mcmc_writers.save_mcmc_results(result, output_dir, prefix='mcmc')[source]

Save MCMC/CMC results to files.

Creates: - {prefix}_summary.json: Parameter summaries with credible intervals - {prefix}_diagnostics.json: Convergence diagnostics (R-hat, ESS) - {prefix}_samples.npz: Full posterior samples (compressed)

Parameters:

  • result (CMCResult) – CMC result object

  • output_dir (Path | str) – Output directory

  • prefix (str) – Filename prefix

Return type:

dict[str, Path]

Returns:

Dict mapping file type to saved path

heterodyne.io.mcmc_writers.save_mcmc_diagnostics(result, output_path, r_hat_threshold=1.1, min_bfmi=0.3)[source]

Save MCMC convergence diagnostics.

Parameters:

  • result (CMCResult) – CMC result object

  • output_path (Path | str) – Output file path

  • r_hat_threshold (float) – R-hat convergence threshold (default 1.1)

  • min_bfmi (float) – Minimum BFMI threshold (default 0.3)

Return type:

Path

Returns:

Path to saved file

heterodyne.io.mcmc_writers.format_mcmc_summary(result)[source]

Format MCMC result as human-readable summary.

Parameters:

result (CMCResult) – CMC result object

Return type:

str

Returns:

Formatted summary string