refactor of group 6 (utilities and helpers) completed: improve utilities,...

import numpy as np

import jax

from .checkpoint import save_pickle

from .utils.logging_utils import format_time

from .logging import _format_rate

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

                save_path, random_seed, num_chains, requested_num_chains, num_warmup,

                device_type, sampler_type, param_config, config, format_time,

                _format_rate, log_func, start_offset=0, existing_grouped=None,

                existing_energy=None, track_energy=True, label="sampling",

                skip_jit_pilot=False):

def run_batched(mcmc_obj, target, chain_count, ctx, start_offset=0,

                existing_grouped=None, existing_energy=None, track_energy=True,

                label="sampling", skip_jit_pilot=False):

    """Run MCMC in checkpoint_interval-sized batches with progress logging."""

    plData = ctx.plData

    plModel = ctx.plModel

    checkpoint_interval = ctx.checkpoint_interval

    save_path = ctx.save_path

    random_seed = ctx.random_seed

    num_chains = ctx.num_chains

    requested_num_chains = ctx.requested_num_chains

    num_warmup = ctx.num_warmup

    device_type = ctx.device_type

    sampler_type = ctx.sampler_type

    param_config = ctx.param_config

    config = ctx.config

    log_func = ctx.log_func

    samples_grouped = existing_grouped or {}

    energy_grouped = existing_energy

    samples_so_far = start_offset

    Unified batched-sampling loop extracted from runSampler_NumPyro.

    """

    from .sampler_utils import run_batched as _rb

    from .utils.logging_utils import format_time

    from .logging import _format_rate

    return _rb(mcmc_obj, target, chain_count, ctx.plData, ctx.plModel, ctx.checkpoint_interval,

               ctx.save_path, ctx.random_seed, ctx.num_chains, ctx.requested_num_chains, ctx.num_warmup,

               ctx.device_type, ctx.sampler_type, ctx.param_config, ctx.config, format_time,

               _format_rate, ctx.log_func, start_offset=start_offset,

    return _rb(mcmc_obj, target, chain_count, ctx, start_offset=start_offset,

               existing_grouped=existing_grouped, existing_energy=existing_energy,

               track_energy=track_energy, label=label, skip_jit_pilot=skip_jit_pilot)

    try:

        return jax.tree_util.tree_map(_cast_leaf, state)

    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.

        # If casting fails, fail fast. A dtype mismatch in the warmup state

        # will inevitably cause a crash in the sampler's first step.

        _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

        msg = f"CRITICAL: failed to cast warmup state to {target_dtype}: {e}. " \

              "The checkpoint precision is incompatible with the current session."

        _log(plModel, msg)

        raise RuntimeError(msg) from e

def _fold_times_f32(times, t0, period):

    """Phase-fold times in float64 and return float32 local times near t0=0.

import jax.numpy as jnp

import numpy as np

import math

import ast

def linearizeRollAngle(time, ra, visit, pdf_filename=None, onlyConversion=False):

    """

        'abs': abs,

    }

    # Only allow safe characters

    allowed_chars = set('0123456789.+-*/() eipqrtlogxsncab')

    if not all(c in allowed_chars for c in clean_expr.replace('sqrt', '').replace('log10', '').replace('log', '').replace('exp', '').replace('sin', '').replace('cos', '').replace('tan', '').replace('abs', '').replace('pi', '').replace('inf', '')):

        raise ValueError(f"Unsafe characters in expression: {expr}")

    try:

        result = eval(clean_expr, {"__builtins__": {}}, safe_namespace)

        return float(result)

        node = ast.parse(clean_expr, mode='eval').body

        def _eval(n):

            if isinstance(n, ast.Constant):

                return n.value

            if isinstance(n, ast.BinOp):

                ops = {

                    ast.Add: lambda a, b: a + b,

                    ast.Sub: lambda a, b: a - b,

                    ast.Mult: lambda a, b: a * b,

                    ast.Div: lambda a, b: a / b,

                    ast.Pow: lambda a, b: a ** b,

                }

                return ops[type(n.op)](_eval(n.left), _eval(n.right))

            if isinstance(n, ast.UnaryOp):

                ops = {

                    ast.UAdd: lambda a: +a,

                    ast.USub: lambda a: -a,

                }

                return ops[type(n.op)](_eval(n.operand))

            if isinstance(n, ast.Call):

                if not isinstance(n.func, ast.Name):

                    raise ValueError('Only simple function calls allowed')

                func_name = n.func.id

                if func_name not in safe_namespace:

                    raise ValueError(f'Unsupported function: {func_name}')

                args = [_eval(arg) for arg in n.args]

                return safe_namespace[func_name](*args)

            if isinstance(n, ast.Name):

                if n.id not in safe_namespace:

                    raise ValueError(f'Unsupported name: {n.id}')

                return safe_namespace[n.id]

            raise ValueError(f'Unsupported AST node: {type(n)}')

        return float(_eval(node))

    except Exception as e:

        raise ValueError(f"Cannot evaluate expression '{expr}': {e}")