Skip to content

theta 分 bin 改进及 jax 版本的重建

Liu Xuewei requested to merge thetauniform into master

/cc @liuyq @heroxbd

需进一步检查 recon_jax.py 的代码实现及重建结果。

之前关于 interpax.CublicSpline 的问题解决了，代码在 recon_jax_CubicSpline.py (临时测试没有 push 上来）。

报错信息：

{ time  ./recon_jax.py evtsel/1319/001.h5 -o recon_jax/1319/001.h5 \
        --sipm-params sipmpars/SiPMParams.h5 \
        --time-offset sipmpars/timeoffset.npz \
        --response svd/r_50_3.npz \
        --geom sipmpars/channelpos.txt; }
[*] Initialized JAX Backend: CPU
[*] Loading Response Model: svd/r_50_3.npz
[*] Detected 694 Bad Channels (rPDE < 0).
[*] JIT Compiler Warmup...
Traceback (most recent call last):
  File "/mnt/stage/xuewei/work1/TAO/TAORec_diff/./recon_jax.py", line 567, in <module>
    main()
  File "/mnt/stage/xuewei/work1/TAO/TAORec_diff/./recon_jax.py", line 374, in main
    nll_val_and_grad(jnp.array([0.0, 0.0, 0.0, 1.0]), dummy_Q, dummy_fired)
  File "/mnt/stage/xuewei/work1/TAO/TAORec_diff/./recon_jax.py", line 187, in nll_static
    lam0_all = fast_predict(x, y, z) * rpde_arr_j
               ^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/stage/xuewei/work1/TAO/TAORec_diff/./recon_jax.py", line 138, in fast_predict
    w = jax.vmap(calc_w, in_axes=1)(W_k_j)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/stage/xuewei/work1/TAO/TAORec_diff/./recon_jax.py", line 136, in calc_w
    return interpax.CubicSpline(r_axis_j, w_arr, bc_type="clamped")(r)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/xuewei/.pyvenv/lib/python3.11/site-packages/equinox/_module/_module.py", line 396, in __call__
    self = tryself = super().__call__(*args, **kwargs)  # pyright: ignore[reportAttributeAccessIssue]
    ^^^^^^^^^^^
  File "/home/xuewei/.pyvenv/lib/python3.11/site-packages/equinox/_module/_better_abstract.py", line 233, in __call__
    self = super().__call__(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/xuewei/.pyvenv/lib/python3.11/site-packages/interpax/_ppoly.py", line 807, in __init__
    x, _, y, axis, _ = prepare_input(x, y, axis, check=check)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/xuewei/.pyvenv/lib/python3.11/site-packages/interpax/_ppoly.py", line 481, in prepare_input
    not jnp.all(jnp.isfinite(x)),
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
jax.errors.TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].
The error occurred while tracing the function fast_predict at /mnt/stage/xuewei/work1/TAO/TAORec_diff/./recon_jax.py:115 for jit. This value became a tracer due to JAX operations on these lines:

  operation a:f64[] = max b 1e-10:f64[]
    from line /mnt/stage/xuewei/work1/TAO/TAORec_diff/./recon_jax.py:119:12 (build_jax_functions.<locals>.fast_predict)

  operation a:f64[] = mul b 939.584:f64[]
    from line /mnt/stage/xuewei/work1/TAO/TAORec_diff/./recon_jax.py:121:47 (build_jax_functions.<locals>.fast_predict)

  operation b:f64[8048] = jit[
  name=clip
  jaxpr={ lambda ; a:f64[8048] c:f64[] d:f64[]. let
      e:f64[] = convert_element_type[new_dtype=float64 weak_type=False] c
      f:f64[8048] = max e a
      g:f64[] = convert_element_type[new_dtype=float64 weak_type=False] d
      b:f64[8048] = min g f
    in (b,) }
] a -1.0:f64[] 1.0:f64[]
    from line /mnt/stage/xuewei/work1/TAO/TAORec_diff/./recon_jax.py:122:19 (build_jax_functions.<locals>.fast_predict)

  operation a:f64[] = add b 882818.093056:f64[]
    from line /mnt/stage/xuewei/work1/TAO/TAORec_diff/./recon_jax.py:127:15 (build_jax_functions.<locals>.fast_predict)

  operation a:f64[] = mul 2.0:f64[] b
    from line /mnt/stage/xuewei/work1/TAO/TAORec_diff/./recon_jax.py:127:36 (build_jax_functions.<locals>.fast_predict)

(Additional originating lines are not shown.)
See https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerBoolConversionError
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

需要把 interpax.CubicSpline的构建过程提到了 @jax.jit 的外面（预编译阶段）。新的 recon_jax.py 实现没有这个问题。

Merge request reports