refactor of group 6 (utilities and helpers) started

                            existing_samples[name] = arr.astype(np.float64)

                    log_func(logFile, f"  Upcast existing samples to float64")

            else:

                existing_warmup_state = _cast_warmup_state_dtype(existing_warmup_state, use_float64=False)

                existing_warmup_state = _cast_warmup_state_dtype(

                    existing_warmup_state, use_float64=False, log_func=log_func, plModel=logFile)

            log_func(logFile, "")

        except Exception as e:

import time as ttt

import numpy as np

import jax

import jax.numpy as jnp

import numpyro

from .checkpoint import save_pickle

def run_batched(mcmc_obj, target, chain_count, plData, plModel, checkpoint_interval,

        else:

            moves = {ESS.DifferentialMove(): 0.7, ESS.KDEMove(): 0.3}

        # return ESS(model_fn, moves=moves, init_strategy=init_to_uniform(radius=ensemble_init_radius))

        from numpyro.infer import init_to_median

        return ESS(model_fn, moves=moves, init_strategy=init_to_median())

import os

"""Small shared logging/formatting helpers used across the sampling engine."""

import time as ttt

def _log(log_file, msg):

    """Write message to log file."""

    """Append a timestamped message to a log file.

    Args:

        log_file: Either a path (str) to the log file, or an object exposing

            a ``logFile`` attribute (e.g. a model object) pointing to the path.

        msg (str): The message to write.

    """

    # Handle case where log_file might be a model object with a logFile attribute

    path = log_file.logFile if hasattr(log_file, 'logFile') else log_file

    with open(path, 'a') as flog:

        flog.write(ttt.strftime("%Y-%m-%d %H:%M:%S", ttt.gmtime()) + ' ' + msg + '\n')

def format_time(seconds):

    """Format a duration in seconds as a human-readable string (s / min / h).

    Args:

        seconds (float): Duration in seconds.

    Returns:

        str: Formatted duration, e.g. ``"45s"``, ``"3.2min"``, ``"1.5h"``.

    """

    if seconds < 60:

        return f"{seconds:.0f}s"

    elif seconds < 3600:

    is_recommended_x64 = bool(precision_critical or ref_delta_planets or gp_configs)

    if use_float64 == True or use_float64 == 'true':

    if use_float64 is True or use_float64 == 'true':

        jax.config.update('jax_enable_x64', True)

        _log(plModel, "JAX precision: float64 (manual override via mcmc.use_float64)")

        return True

    elif use_float64 == False or use_float64 == 'false':

    elif use_float64 is False or use_float64 == 'false':

        jax.config.update('jax_enable_x64', False)

        if is_recommended_x64:

            warn_msg = (

        return False

def _cast_warmup_state_dtype(state, use_float64):

def _cast_warmup_state_dtype(state, use_float64, log_func=None, plModel=None):

    """Cast all floating-point arrays in an ESS/AIES warmup state to match the current JAX

    precision setting.

    Args:

        state: ESS/AIES MCMC last_state (any JAX pytree).

        use_float64: if True, cast to float64; if False, cast to float32.

        log_func: Optional logging function with signature ``log_func(plModel, msg)``,

            used to report a casting failure. Falls back to ``print`` if not provided.

        plModel: Optional model object forwarded to ``log_func`` (see ``_configure_jax_precision``).

    Returns:

        State with all floating-point leaves recast, or the original state if casting fails.

    except Exception as e:

        # If casting fails for any reason, return as-is and let the sampler

        # surface the original dtype error with a clear traceback.

        _log = log_func if log_func is not None else print

        _log(plModel, f"WARNING: failed to cast warmup state to {target_dtype}: {e}. "

                       "Returning state unchanged; a dtype-mismatch error may surface later.")

        return state

def _fold_times_f32(times, t0, period):

import jax

import jax.numpy as jnp

import numpy as np

import math

def relabel_array(arr):

    # 'inverse' returns an array of the same shape as 'arr'

    # where each element is the index of that value in the unique array

    # Use a static size for jnp.unique to avoid dynamic-output-shape bugs

    # in JAX's Thunk runtime on newer XLA CPU backends.

    """

    Relabel an array's values with 1-based sequential integer labels.

    'inverse' returns an array of the same shape as 'arr' where each element

    is the index of that value in the unique array. Uses a static size for

    jnp.unique to avoid dynamic-output-shape bugs in JAX's Thunk runtime on

    newer XLA CPU backends.

    Args:

        arr: Array-like of (typically categorical) values to relabel.

    Returns:

        tuple: (new_labels, unique_vals) where new_labels has the same shape

        as ``arr`` with values replaced by their 1-based index in the sorted

        unique values, and unique_vals is the array of unique values found.

    """

    unique_vals, inverse = jnp.unique(

        arr,

        return_inverse=True,

    Raises:

        ValueError: If expression cannot be safely evaluated

    """

    import math

    # Clean up common prefixes

    clean_expr = expr.replace('jnp.', '').replace('np.', '').replace('math.', '')