+ 1import types
+ 2import warnings
+ 3
+ 4import numpy as np
+ 5
+ 6import pyxu.abc as pxa
+ 7import pyxu.info.deps as pxd
+ 8import pyxu.info.ptype as pxt
+ 9import pyxu.info.warning as pxw
+ 10import pyxu.runtime as pxrt
+ 11import pyxu.util as pxu
+ 12
+ 13
+
+[docs]
+ 14class RayXRT(pxa.LinOp):
+ 15 r"""
+ 16 2D/3D X-Ray Transform :math:`\mathcal{P}[f](\mathbf{n}, \mathbf{t})`.
+ 17
+ 18 This implementation computes the XRT using a ray-marching method based on the `Dr.Jit
+ 19 <https://drjit.readthedocs.io/en/latest/reference.html>`_ compiler.
+ 20
+ 21 See :py:class:`~pyxu.experimental.xray.XRayTransform` for notational conventions used herein.
+ 22 """
+ 23
+
+[docs]
+ 24 def __init__(
+ 25 self,
+ 26 arg_shape,
+ 27 origin,
+ 28 pitch,
+ 29 n_spec,
+ 30 t_spec,
+ 31 enable_warnings: bool = True,
+ 32 ):
+ 33 r"""
+ 34 Parameters
+ 35 ----------
+ 36 arg_shape: NDArrayShape
+ 37 Pixel count in each dimension.
+ 38 origin: NDArray
+ 39 Bottom-left coordinate :math:`\mathbf{o} \in \mathbb{R}^{D}`.
+ 40 pitch: NDArray
+ 41 Pixel size :math:`\mathbf{\Delta} \in \mathbb{R}_{+}^{D}`.
+ 42 n_spec: NDArray
+ 43 (N_l, D) ray directions :math:`\mathbf{n} \in \mathbb{S}^{D-1}`.
+ 44 t_spec: NDArray
+ 45 (N_l, D) offset specifiers :math:`\mathbf{t} \in \mathbb{R}^{D}`.
+ 46 enable_warnings: bool
+ 47
+ 48
+ 49 .. warning::
+ 50
+ 51 This is a low-level interface which does not perform any parameter validation. Users are expected to
+ 52 instantiate :py:class:`~pyxu.experimental.xray._rt.RayXRT` by calling
+ 53 :py:meth:`~pyxu.experimental.xray.XRayTransform.init` instead.
+ 54 """
+ 55 N_px = np.prod(arg_shape)
+ 56 N_l = len(n_spec)
+ 57 super().__init__(shape=(N_l, N_px))
+ 58
+ 59 W = pxrt.Width # shorthand
+ 60 n_spec = n_spec.astype(W.SINGLE.value, copy=False)
+ 61 t_spec = t_spec.astype(W.SINGLE.value, copy=False)
+ 62
+ 63 # Internal variables. (Have shapes consistent with user inputs.)
+ 64 self._arg_shape = arg_shape # (D,)
+ 65 self._origin = origin # (D,)
+ 66 self._pitch = pitch # (D,)
+ 67 self._ray_n = n_spec # (N_l, D)
+ 68 self._ray_t = t_spec # (N_l, D)
+ 69 self._ndi = pxd.NDArrayInfo.from_obj(n_spec)
+ 70 self._enable_warnings = bool(enable_warnings)
+ 71
+ 72 # Cheap analytical Lipschitz upper bound given by
+ 73 # \sigma_{\max}(P) <= \max{P.sum(axis=0), P.sum(axis=1)},
+ 74 # with
+ 75 # P.sum(axis=0) <= \norm{pitch}{2} * N_ray [very pessimistic]
+ 76 # P.sum(axis=1) <= \norm{pitch}{2} * \norm{arg_shape}{2}
+ 77 row_ub = np.linalg.norm(pitch) * np.linalg.norm(arg_shape)
+ 78 col_ub = np.linalg.norm(pitch) * len(n_spec)
+ 79 self.lipschitz = max(row_ub, col_ub)
+ 80
+ 81 # Dr.Jit variables. {Have shapes consistent for xrt_[apply,adjoint]().}
+ 82 # xrt_[apply,adjoint]() only support D=3 case.
+ 83 # -> D=2 case is embedded into 3D.
+ 84 drb = _load_dr_variant(self._ndi)
+ 85 Ray3f, _, self._fwd, self._bwd = _get_dr_obj(self._ndi)
+ 86 if len(arg_shape) == 2:
+ 87 self._dr = dict(
+ 88 o=drb.Array3f(*self._origin, 0),
+ 89 pitch=drb.Array3f(*self._pitch, 1),
+ 90 N=drb.Array3u(*self._arg_shape, 1),
+ 91 r=Ray3f(
+ 92 o=drb.Array3f(*[_xp2dr(_) for _ in self._ray_t.T], 0.5), # Z-extension mid-point
+ 93 d=drb.Array3f(*[_xp2dr(_) for _ in self._ray_n.T], 0),
+ 94 ),
+ 95 )
+ 96 else:
+ 97 self._dr = dict(
+ 98 o=drb.Array3f(*self._origin),
+ 99 pitch=drb.Array3f(*self._pitch),
+100 N=drb.Array3u(*self._arg_shape),
+101 r=Ray3f(
+102 o=drb.Array3f(*[_xp2dr(_) for _ in self._ray_t.T]),
+103 d=drb.Array3f(*[_xp2dr(_) for _ in self._ray_n.T]),
+104 ),
+105 )
+
+106
+
+[docs]
+107 @pxrt.enforce_precision(i="arr")
+108 @pxu.vectorize(i="arr")
+109 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
+110 r"""
+111 Parameters
+112 ----------
+113 arr: NDArray
+114 (..., arg_shape.prod()) spatial weights.
+115
+116 Returns
+117 -------
+118 P: NDArray
+119 (..., N_l) XRT samples.
+120 """
+121 arr, dtype = self._warn_cast(arr)
+122
+123 drb = _load_dr_variant(self._ndi)
+124 P = self._fwd(
+125 **self._dr,
+126 I=drb.Float(_xp2dr(arr)),
+127 )
+128
+129 xp = self._ndi.module()
+130 P = xp.asarray(P, dtype=dtype)
+131 return P
+
+132
+
+[docs]
+133 @pxrt.enforce_precision(i="arr")
+134 @pxu.vectorize(i="arr")
+135 def adjoint(self, arr: pxt.NDArray) -> pxt.NDArray:
+136 r"""
+137 Parameters
+138 ----------
+139 arr: NDArray
+140 (..., N_l) XRT samples.
+141
+142 Returns
+143 -------
+144 I: NDArray
+145 (..., arg_shape.prod()) spatial weights.
+146 """
+147 arr, dtype = self._warn_cast(arr)
+148
+149 drb = _load_dr_variant(self._ndi)
+150 I = self._bwd( # noqa: E741
+151 **self._dr,
+152 P=drb.Float(_xp2dr(arr)),
+153 )
+154
+155 xp = self._ndi.module()
+156 I = xp.asarray(I, dtype=dtype) # noqa: E741
+157 return I
+
+158
+159 def _warn_cast(self, arr: pxt.NDArray) -> tuple[pxt.NDArray, pxt.DType]:
+160 W = pxrt.Width # shorthand
+161 if W(arr.dtype) != W.SINGLE:
+162 if self._enable_warnings:
+163 msg = f"Only {W.SINGLE}-precision inputs are supported: casting."
+164 warnings.warn(msg, pxw.PrecisionWarning)
+165 out = arr.astype(dtype=W.SINGLE.value)
+166 else:
+167 out = arr
+168 return out, arr.dtype
+169
+170 def asarray(self, **kwargs) -> pxt.NDArray:
+171 # DASK not yet supported, hence we support `xp=DASK` case sub-optimally for now.
+172 xp = kwargs.get("xp", pxd.NDArrayInfo.default().module())
+173 dtype = kwargs.get("dtype", pxrt.getPrecision().value)
+174
+175 # compute array representation using instance's backend.
+176 kwargs.update(xp=self._ndi.module())
+177 A = super().asarray(**kwargs)
+178
+179 # then cast to user specs.
+180 A = xp.array(pxu.to_NUMPY(A), dtype=dtype)
+181 return A
+182
+183 def gram(self) -> pxt.OpT:
+184 # We replace apply() with a variant which does not force evaluation between FW/BW calls.
+185
+186 @pxrt.enforce_precision(i="arr")
+187 @pxu.vectorize(i="arr")
+188 def op_apply(_, arr: pxt.NDArray) -> pxt.NDArray:
+189 arr, dtype = _._warn_cast(arr)
+190
+191 drb = _load_dr_variant(_._ndi)
+192 P = _._fwd(
+193 **_._dr,
+194 I=drb.Float(_xp2dr(arr)),
+195 )
+196 I = _._bwd( # noqa: E741
+197 **_._dr,
+198 P=P,
+199 )
+200
+201 xp = _._ndi.module()
+202 I = xp.asarray(I, dtype=dtype) # noqa: E741
+203 return I
+204
+205 op = super().gram()
+206 op.apply = types.MethodType(op_apply, self)
+207 return op
+208
+209 def cogram(self) -> pxt.OpT:
+210 # We replace apply() with a variant which does not enforce evaluation between BW/FW calls.
+211
+212 @pxrt.enforce_precision(i="arr")
+213 @pxu.vectorize(i="arr")
+214 def op_apply(_, arr: pxt.NDArray) -> pxt.NDArray:
+215 arr, dtype = _._warn_cast(arr)
+216
+217 drb = _load_dr_variant(_._ndi)
+218 I = _._bwd( # noqa: E741
+219 **_._dr,
+220 P=drb.Float(_xp2dr(arr)),
+221 )
+222 P = _._fwd(
+223 **_._dr,
+224 I=I,
+225 )
+226
+227 xp = _._ndi.module()
+228 P = xp.asarray(P, dtype=dtype)
+229 return P
+230
+231 op = super().cogram()
+232 op.apply = types.MethodType(op_apply, self)
+233 return op
+234
+
+[docs]
+235 def diagnostic_plot(self, ray_idx: pxt.NDArray = None):
+236 r"""
+237 Plot ray trajectories.
+238
+239 Parameters
+240 ----------
+241 ray_idx: NDArray
+242 (Q,) indices of rays to plot. (Default: show all rays.)
+243
+244 Returns
+245 -------
+246 fig: :py:class:`~matplotlib.figure.Figure`
+247 Diagnostic plot.
+248
+249 Notes
+250 -----
+251 * Rays which do not intersect the volume are **not** shown.
+252
+253 Examples
+254 --------
+255
+256 .. plot::
+257
+258 import numpy as np
+259 import pyxu.experimental.xray as pxr
+260
+261 op = pxr.XRayTransform.init(
+262 arg_shape=(5, 6),
+263 origin=0,
+264 pitch=1,
+265 method="ray-trace",
+266 n_spec=np.array([[1 , 0 ], # 3 rays ...
+267 [0.5 , 0.5 ],
+268 [0.75, 0.25]]),
+269 t_spec=np.array([[2.5, 3], # ... all defined w.r.t volume center
+270 [2.5, 3],
+271 [2.5, 3]]),
+272 )
+273 fig = op.diagnostic_plot()
+274 fig.show()
+275
+276 Notes
+277 -----
+278 Requires `Matplotlib <https://matplotlib.org/>`_ to be installed.
+279 """
+280 dr = pxu.import_module("drjit")
+281 plt = pxu.import_module("matplotlib.pyplot")
+282 collections = pxu.import_module("matplotlib.collections")
+283 patches = pxu.import_module("matplotlib.patches")
+284
+285 # Setup Figure ========================================================
+286 D = len(self._arg_shape)
+287 if D == 2:
+288 fig, ax = plt.subplots()
+289 data = [(ax, [0, 1], ["x", "y"])]
+290 else: # D == 3 case
+291 fig, ax = plt.subplots(ncols=3)
+292 data = [
+293 (ax[0], [0, 1], ["x", "y"]),
+294 (ax[1], [0, 2], ["x", "z"]),
+295 (ax[2], [1, 2], ["y", "z"]),
+296 ]
+297
+298 # Determine which rays intersect with BoundingBox =====================
+299 _, BBox3f, _, _ = _get_dr_obj(self._ndi)
+300 active, a1, a2 = BBox3f(
+301 pMin=self._dr["o"],
+302 pMax=self._dr["o"] + self._dr["pitch"] * self._dr["N"],
+303 ).ray_intersect(self._dr["r"])
+304 dr.eval(active, a1, a2)
+305
+306 # Then extract subset of interest (which intersect bbox)
+307 if ray_idx is None:
+308 ray_idx = slice(None)
+309 active, a1, a2 = map(lambda _: _.numpy()[ray_idx], [active, a1, a2]) # (Q,)
+310 a12 = np.stack([a1, a2], axis=-1)[active] # (N_active, 2)
+311 ray_n = self._ray_n[ray_idx][active] # (N_active, D)
+312 ray_t = self._ray_t[ray_idx][active] # (N_active, D)
+313
+314 for _ax, dim_idx, dim_label in data:
+315 # Subsample right dimensions ======================================
+316 origin = np.array(self._origin)[dim_idx]
+317 arg_shape = np.array(self._arg_shape)[dim_idx]
+318 pitch = np.array(self._pitch)[dim_idx]
+319 _ray_n = ray_n[:, dim_idx]
+320 _ray_t = ray_t[:, dim_idx]
+321
+322 # Helper variables ================================================
+323 bbox_dim = arg_shape * pitch
+324
+325 # Draw BBox =======================================================
+326 rect = patches.Rectangle(
+327 xy=origin,
+328 width=bbox_dim[0],
+329 height=bbox_dim[1],
+330 facecolor="none",
+331 edgecolor="k",
+332 label="volume BBox",
+333 )
+334 _ax.add_patch(rect)
+335
+336 # Draw Pitch =======================================================
+337 p_rect = patches.Rectangle(
+338 xy=origin + bbox_dim - pitch,
+339 width=pitch[0],
+340 height=pitch[1],
+341 facecolor="r",
+342 edgecolor="none",
+343 label="pitch size",
+344 )
+345 _ax.add_patch(p_rect)
+346
+347 # Draw Origin =====================================================
+348 _ax.scatter(
+349 origin[0],
+350 origin[1],
+351 color="k",
+352 label="origin",
+353 marker="x",
+354 )
+355
+356 # Draw Rays & Anchor Points =======================================
+357 # Each (2,2) sub-array in `coords` represents line start/end coordinates.
+358 coords = _ray_t.reshape(-1, 1, 2) + a12.reshape(-1, 2, 1) * _ray_n.reshape(-1, 1, 2) # (N_active, 2, 2)
+359 lines = collections.LineCollection(
+360 coords,
+361 label=r"$t + \alpha n$",
+362 color="k",
+363 alpha=0.5,
+364 linewidth=1,
+365 )
+366 _ax.add_collection(lines)
+367 _ax.scatter(
+368 _ray_t[:, 0],
+369 _ray_t[:, 1],
+370 label=r"t",
+371 color="g",
+372 marker=".",
+373 )
+374
+375 # Misc Details ====================================================
+376 pad_width = 0.1 * bbox_dim # 10% axial pad
+377 _ax.set_xlabel(dim_label[0])
+378 _ax.set_ylabel(dim_label[1])
+379 _ax.set_xlim(origin[0] - pad_width[0], origin[0] + bbox_dim[0] + pad_width[0])
+380 _ax.set_ylim(origin[1] - pad_width[1], origin[1] + bbox_dim[1] + pad_width[1])
+381 _ax.legend(loc="lower right", bbox_to_anchor=(1, 1))
+382 _ax.set_aspect(1)
+383
+384 fig.tight_layout()
+385 return fig
+
+
+386
+387
+388def _xp2dr(x: pxt.NDArray):
+389 # Transform NP/CP inputs to format allowing zero-copy casts to {drb.Float, drb.Array3f}
+390 ndi = pxd.NDArrayInfo.from_obj(x)
+391 if ndi == pxd.NDArrayInfo.NUMPY:
+392 return x
+393 elif ndi == pxd.NDArrayInfo.CUPY:
+394 return x.__dlpack__()
+395 else:
+396 raise NotImplementedError
+397
+398
+399def _load_dr_variant(ndi: pxd.NDArrayInfo):
+400 if ndi == pxd.NDArrayInfo.NUMPY:
+401 drb = pxu.import_module("drjit.llvm")
+402 elif ndi == pxd.NDArrayInfo.CUPY:
+403 drb = pxu.import_module("drjit.cuda")
+404 else:
+405 raise NotImplementedError
+406 return drb
+407
+408
+409def _get_dr_obj(ndi: pxd.NDArrayInfo):
+410 # Create DrJIT objects needed for fwd/bwd transforms.
+411 import drjit as dr
+412
+413 drb = _load_dr_variant(ndi)
+414
+415 class Ray3f:
+416 # Dr.JIT-backed 3D ray.
+417 #
+418 # Rays take the parametric form
+419 # r(t) = o + d * t,
+420 # where {o, d} \in \bR^{3} and t \in \bR.
+421 DRJIT_STRUCT = dict(
+422 o=drb.Array3f,
+423 d=drb.Array3f,
+424 )
+425
+426 def __init__(self, o=drb.Array3f(), d=drb.Array3f()):
+427 # Parameters
+428 # ----------
+429 # o: Array3f
+430 # (3,) ray origin.
+431 # d: Array3f
+432 # (3,) ray direction.
+433 #
+434 # [2023.10.19 Sepand/Miguel]
+435 # Use C++'s `operator=()` semantics instead of Python's `=` to safely reference inputs.
+436 self.o = drb.Array3f()
+437 self.o.assign(o)
+438
+439 self.d = drb.Array3f()
+440 self.d.assign(d)
+441
+442 def __call__(self, t: drb.Float) -> drb.Array3f:
+443 # Compute r(t).
+444 #
+445 # Parameters
+446 # ----------
+447 # t: Float
+448 #
+449 # Returns
+450 # -------
+451 # p: Array3f
+452 # p = r(t)
+453 return dr.fma(self.d, t, self.o)
+454
+455 def assign(self, r: "Ray3f"):
+456 # See __init__'s docstring for more info.
+457 self.o.assign(r.o)
+458 self.d.assign(r.d)
+459
+460 def __repr__(self) -> str:
+461 return f"Ray3f(o={dr.shape(self.o)}, d={dr.shape(self.d)})"
+462
+463 class BoundingBox3f:
+464 # Dr.JIT-backed 3D bounding box.
+465 #
+466 # A bounding box is described by coordinates {pMin, pMax} of two of its diagonal corners.
+467 #
+468 # Important
+469 # ---------
+470 # We do NOT check if (pMin < pMax) when constructing the BBox: users are left responsible of their actions.
+471 DRJIT_STRUCT = dict(
+472 pMin=drb.Array3f,
+473 pMax=drb.Array3f,
+474 )
+475
+476 def __init__(self, pMin=drb.Array3f(), pMax=drb.Array3f()):
+477 # Parameters
+478 # ----------
+479 # pMin: Array3f
+480 # (3,) corner coordinate.
+481 # pMax: Array3f
+482 # (3,) corner coordinate.
+483 #
+484 # [2023.10.19 Sepand/Miguel]
+485 # Use C++'s `operator=()` semantics instead of Python's `=` to safely reference inputs.
+486 self.pMin = drb.Array3f()
+487 self.pMin.assign(pMin)
+488
+489 self.pMax = drb.Array3f()
+490 self.pMax.assign(pMax)
+491
+492 def contains(self, p: drb.Array3f) -> drb.Bool:
+493 # Check if point `p` lies in/on the BBox.
+494 return dr.all((self.pMin <= p) & (p <= self.pMax))
+495
+496 def ray_intersect(self, r: Ray3f) -> tuple[drb.Bool, drb.Float, drb.Float]:
+497 # Compute ray/bbox intersection parameters. [Adapted from Mitsuba3's ray_intersect().]
+498 #
+499 # Parameters
+500 # ----------
+501 # r: Ray3f
+502 #
+503 # Returns
+504 # -------
+505 # active: Bool
+506 # True if intersection occurs.
+507 # mint, maxt: Float
+508 # Ray parameters `t` such that r(t) touches a BBox border.
+509 # The value only makes sense if `active` is enabled.
+510
+511 # Ensure ray has a nonzero slope on each axis, or that its origin on a 0-valued axis is within the box bounds.
+512 active = dr.all(dr.neq(r.d, 0) | (self.pMin < r.o) | (r.o < self.pMax))
+513
+514 # Compute intersection intervals for each axis
+515 d_rcp = dr.rcp(r.d)
+516 t1 = (self.pMin - r.o) * d_rcp
+517 t2 = (self.pMax - r.o) * d_rcp
+518
+519 # Ensure proper ordering
+520 t1p = dr.minimum(t1, t2)
+521 t2p = dr.maximum(t1, t2)
+522
+523 # Intersect intervals
+524 mint = dr.max(t1p)
+525 maxt = dr.min(t2p)
+526 active &= mint <= maxt
+527
+528 return active, mint, maxt
+529
+530 def assign(self, bbox: "BoundingBox3f"):
+531 # See __init__'s docstring for more info.
+532 self.pMin.assign(bbox.pMin)
+533 self.pMax.assign(bbox.pMax)
+534
+535 def __repr__(self) -> str:
+536 return f"BoundingBox3f(pMin={dr.shape(self.pMin)}, pMax={dr.shape(self.pMax)})"
+537
+538 def ray_step(r: Ray3f) -> Ray3f:
+539 # Advance ray until next unit-step lattice intersection.
+540 #
+541 # Parameters
+542 # r(o, d): ray to move. (`d` assumed normalized.)
+543 # Returns
+544 # r_next(o_next, d): next ray position on unit-step lattice intersection.
+545 eps = 1e-4 # threshold for proximity tests with 0
+546
+547 # Go to local coordinate system.
+548 o_ref = dr.floor(r.o)
+549 r_local = Ray3f(o=r.o - o_ref, d=r.d)
+550
+551 # Find bounding box for ray-intersection tests.
+552 on_boundary = r_local.o <= eps
+553 bbox_border = dr.select(on_boundary, dr.sign(r.d), 1)
+554 bbox = BoundingBox3f(
+555 dr.minimum(0, bbox_border),
+556 dr.maximum(0, bbox_border),
+557 )
+558
+559 # Compute step size to closest bounding box wall.
+560 # (a1, a2) may contain negative values or Infs.
+561 # In any case, we must always choose min(a1, a2) > 0.
+562 _, a1, a2 = bbox.ray_intersect(r_local)
+563 a_min = dr.minimum(a1, a2)
+564 a_max = dr.maximum(a1, a2)
+565 a = dr.select(a_min >= eps, a_min, a_max)
+566
+567 # Move ray to new position in global coordinates.
+568 # r_next located on lattice intersection (up to FP error).
+569 r_next = Ray3f(o=o_ref + r_local(a), d=r.d)
+570 return r_next
+571
+572 def xrt_apply(
+573 o: drb.Array3f,
+574 pitch: drb.Array3f,
+575 N: drb.Array3u,
+576 I: drb.Float,
+577 r: Ray3f,
+578 ) -> drb.Float:
+579 # X-Ray Forward-Projection.
+580 #
+581 # Parameters
+582 # o: bottom-left coordinate of I[0,0,0]
+583 # pitch: cell dimensions \in \bR_{+}
+584 # N: (Nx,Ny,Nz) lattice size
+585 # I: (Nx*Ny*Nz,) cell weights \in \bR [C-order]
+586 # r: (L,) ray descriptors
+587 # Returns
+588 # P: (L,) forward-projected samples \in \bR
+589
+590 # Go to normalized coordinates
+591 ipitch = dr.rcp(pitch)
+592 r = Ray3f(
+593 o=(r.o - o) * ipitch,
+594 d=dr.normalize(r.d * ipitch),
+595 )
+596 stride = drb.Array3u(N[1] * N[2], N[2], 1)
+597 flat_index = lambda i: dr.dot(stride, drb.Array3u(i)) # Array3f (int-valued) -> UInt32
+598
+599 L = max(dr.shape(r.o)[1], dr.shape(r.d)[1])
+600 P = dr.zeros(drb.Float, shape=L) # Forward-Projection samples
+601 idx_P = dr.arange(drb.UInt32, L)
+602
+603 # Move (intersecting) rays to volume surface
+604 bbox_vol = BoundingBox3f(drb.Array3f(0), drb.Array3f(N))
+605 active, a1, a2 = bbox_vol.ray_intersect(r)
+606 a_min = dr.minimum(a1, a2)
+607 r.o.assign(dr.select(active, r(a_min), r.o))
+608
+609 r_next = ray_step(r)
+610 active &= bbox_vol.contains(r_next.o)
+611 loop = drb.Loop("XRT FW", lambda: (r, r_next, active))
+612 while loop(active):
+613 length = dr.norm((r_next.o - r.o) * pitch)
+614
+615 idx_I = dr.floor(0.5 * (r_next.o + r.o))
+616 weight = dr.gather(
+617 type(I),
+618 I,
+619 flat_index(idx_I),
+620 active & dr.all(idx_I >= 0),
+621 # Careful to disable out-of-bound queries.
+622 # [This may occur if FP-error caused r_next(above) to not enter the lattice; auto-rectified at next iteration.]
+623 )
+624
+625 # Update line integral estimates
+626 dr.scatter_reduce(dr.ReduceOp.Add, P, weight * length, idx_P, active)
+627
+628 # Walk to next lattice intersection.
+629 r.assign(r_next)
+630 r_next.assign(ray_step(r))
+631 active &= bbox_vol.contains(r_next.o)
+632 return P
+633
+634 def xrt_adjoint(
+635 o: drb.Array3f,
+636 pitch: drb.Array3f,
+637 N: drb.Array3u,
+638 P: drb.Float,
+639 r: Ray3f,
+640 ) -> drb.Float:
+641 # X-Ray Back-Projection.
+642 #
+643 # Parameters
+644 # o: bottom-left coordinate of I[0,0,0]
+645 # pitch: cell dimensions \in \bR_{+}
+646 # N: (Nx,Ny,Nz) lattice size
+647 # P: (L,) X-Ray samples \in \bR
+648 # r: (L,) ray descriptors
+649 # Returns
+650 # I: (Nx*Ny*Nz,) back-projected samples \in \bR [C-order]
+651
+652 # Go to normalized coordinates
+653 ipitch = dr.rcp(pitch)
+654 r = Ray3f(
+655 o=(r.o - o) * ipitch,
+656 d=dr.normalize(r.d * ipitch),
+657 )
+658 stride = drb.Array3u(N[1] * N[2], N[2], 1)
+659 flat_index = lambda i: dr.dot(stride, drb.Array3u(i)) # Array3f (int-valued) -> UInt32
+660
+661 I = dr.zeros(drb.Float, dr.prod(N)[0]) # noqa: E741 (Back-Projection samples)
+662
+663 # Move (intersecting) rays to volume surface
+664 bbox_vol = BoundingBox3f(drb.Array3f(0), drb.Array3f(N))
+665 active, a1, a2 = bbox_vol.ray_intersect(r)
+666 a_min = dr.minimum(a1, a2)
+667 r.o.assign(dr.select(active, r(a_min), r.o))
+668
+669 r_next = ray_step(r)
+670 active &= bbox_vol.contains(r_next.o)
+671 active &= dr.neq(P, 0)
+672 loop = drb.Loop("XRT BW", lambda: (r, r_next, active))
+673 while loop(active):
+674 length = dr.norm((r_next.o - r.o) * pitch)
+675
+676 idx_I = dr.floor(0.5 * (r_next.o + r.o))
+677 dr.scatter_reduce(
+678 dr.ReduceOp.Add,
+679 I,
+680 P * length,
+681 flat_index(idx_I),
+682 active & dr.all(idx_I >= 0),
+683 # Careful to disable out-of-bound queries.
+684 # [This may occur if FP-error caused r_next(above) to not enter the lattice; auto-rectified at next iteration.]
+685 )
+686
+687 # Walk to next lattice intersection.
+688 r.assign(r_next)
+689 r_next.assign(ray_step(r))
+690 active &= bbox_vol.contains(r_next.o)
+691 return I
+692
+693 return Ray3f, BoundingBox3f, xrt_apply, xrt_adjoint
+166 xp = ndi_n.module()
167 n_spec = n_spec / xp.linalg.norm(n_spec, axis=-1, keepdims=True)
168
-169 # Project `t_spec` to lie in n^{\perp}.
-170 # (Not really required; just to have everything in standardized form.)
-171 rng = xp.random.default_rng()
-172 U = rng.standard_normal((N_l, D, D)) # create reference frames n^{\perp}
-173 U[..., 0] = n_spec
-174 U, _ = xp.linalg.qr(U)
-175 s = U[..., 1:].transpose(0, 2, 1) @ t_spec.reshape(N_l, D, 1)
-176 t_spec = (U[..., 1:] @ s).reshape(N_l, D)
+169 ### Project `t_spec` to lie in n^{\perp}.
+170 ### (We don't do this by default since it messes up diagnostic_plot() analysis.)
+171 # rng = xp.random.default_rng()
+172 # U = rng.standard_normal((N_l, D, D)) # create reference frames n^{\perp}
+173 # U[..., 0] = n_spec
+174 # U, _ = xp.linalg.qr(U)
+175 # s = U[..., 1:].transpose(0, 2, 1) @ t_spec.reshape(N_l, D, 1)
+176 # t_spec = (U[..., 1:] @ s).reshape(N_l, D)
177
178 klass = _rt.RayXRT
179 else: # method == "fourier"
diff --git a/_modules/pyxu/experimental/xray/_kernel.html b/_modules/pyxu/operator/map/kernel.html
similarity index 68%
rename from _modules/pyxu/experimental/xray/_kernel.html
rename to _modules/pyxu/operator/map/kernel.html
index e57021c7..092c2767 100644
--- a/_modules/pyxu/experimental/xray/_kernel.html
+++ b/_modules/pyxu/operator/map/kernel.html
@@ -7,7 +7,7 @@
- pyxu.experimental.xray._kernel — Pyxu Documentation
+ pyxu.operator.map.kernel — Pyxu Documentation
@@ -28,7 +28,7 @@
-
+
@@ -56,7 +56,7 @@
-
+
@@ -433,7 +433,7 @@
+ Source code for pyxu.operator.map.kernel
1import types
2
3import numpy as np
@@ -466,320 +466,366 @@ Source code for pyxu.experimental.xray._kernel
13 "Dirac",
14 "FSSPulse",
15 "Box",
- 16 "TruncatedGaussian",
- 17 "KaiserBessel",
- 18]
- 19
+ 16 "Triangle",
+ 17 "TruncatedGaussian",
+ 18 "KaiserBessel",
+ 19]
20
- 21def _get_module(arr: pxt.NDArray):
- 22 N = pxd.NDArrayInfo
- 23 ndi = N.from_obj(arr)
- 24 if ndi == N.NUMPY:
- 25 xp = N.NUMPY.module()
- 26 sps = pxu.import_module("scipy.special")
- 27 elif ndi == N.CUPY:
- 28 xp = N.CUPY.module()
- 29 sps = pxu.import_module("cupyx.scipy.special")
- 30 else:
- 31 raise ValueError(f"Unsupported array type {ndi}.")
- 32 return xp, sps
- 33
+ 21
+ 22def _get_module(arr: pxt.NDArray):
+ 23 N = pxd.NDArrayInfo
+ 24 ndi = N.from_obj(arr)
+ 25 if ndi == N.NUMPY:
+ 26 xp = N.NUMPY.module()
+ 27 sps = pxu.import_module("scipy.special")
+ 28 elif ndi == N.CUPY:
+ 29 xp = N.CUPY.module()
+ 30 sps = pxu.import_module("cupyx.scipy.special")
+ 31 else:
+ 32 raise ValueError(f"Unsupported array type {ndi}.")
+ 33 return xp, sps
34
+ 35
-[docs]
- 35class FSSPulse(pxa.Func):
- 36 r"""
- 37 1D Finite-Support Symmetric function :math:`f: \mathbb{R} \to \mathbb{R}`.
- 38 """
- 39
- 40 def __init__(self):
- 41 super().__init__(shape=(1, 1))
- 42
+[docs]
+ 36class FSSPulse(pxa.Map):
+ 37 r"""
+ 38 1D Finite-Support Symmetric function :math:`f: \mathbb{R} \to \mathbb{R}`, element-wise.
+ 39 """
+ 40
+ 41 def __init__(self, dim: pxt.Integer):
+ 42 super().__init__(shape=(dim, dim))
+ 43
-[docs]
- 43 def support(self) -> pxt.Real:
- 44 r"""
- 45 Returns
- 46 -------
- 47 s: Real
- 48 Value :math:`s > 0` such that :math:`f(x) = 0, \; \forall |x| > s`.
- 49 """
- 50 raise NotImplementedError
-
- 51
+[docs]
+ 44 def support(self) -> pxt.Real:
+ 45 r"""
+ 46 Returns
+ 47 -------
+ 48 s: Real
+ 49 Value :math:`s > 0` such that :math:`f(x) = 0, \; \forall |x| > s`.
+ 50 """
+ 51 raise NotImplementedError
+
+ 52
-[docs]
- 52 def applyF(self, arr: pxt.NDArray) -> pxt.NDArray:
- 53 r"""
- 54 Evaluate :math:`f^{\mathcal{F}}(v)`.
- 55
- 56 :py:meth:`~pyxu.experimental.xray.FSSPulse.applyF` has the same semantics as :py:meth:`~pyxu.abc.Map.apply`.
- 57
- 58 The Fourier convention used is
- 59
- 60 .. math::
- 61
- 62 \mathcal{F}(f)(v) = \int f(x) e^{-j 2\pi v x} dx
- 63 """
- 64 raise NotImplementedError
-
- 65
+[docs]
+ 53 def applyF(self, arr: pxt.NDArray) -> pxt.NDArray:
+ 54 r"""
+ 55 Evaluate :math:`f^{\mathcal{F}}(v)`.
+ 56
+ 57 :py:meth:`~pyxu.operator.FSSPulse.applyF` has the same semantics as :py:meth:`~pyxu.abc.Map.apply`.
+ 58
+ 59 The Fourier convention used is
+ 60
+ 61 .. math::
+ 62
+ 63 \mathcal{F}(f)(v) = \int f(x) e^{-j 2\pi v x} dx
+ 64 """
+ 65 raise NotImplementedError
+
+ 66
-[docs]
- 66 def supportF(self, eps: pxt.Real) -> pxt.Real:
- 67 r"""
- 68 Parameters
- 69 ----------
- 70 eps: Real
- 71 Energy cutoff threshold :math:`\epsilon \in [0, 0.05]`.
- 72
- 73 Returns
- 74 -------
- 75 sF: Real
- 76 Value such that
- 77
- 78 .. math::
- 79
- 80 \int_{-s^{\mathcal{F}}}^{s^{\mathcal{F}}} |f^{\mathcal{F}}(v)|^{2} dv
- 81 \approx
- 82 (1 - \epsilon) \|f\|_{2}^{2}
- 83 """
- 84 eps = float(eps)
- 85 assert 0 <= eps <= 0.05
- 86 tol = 1 - eps
- 87
- 88 def energy(f: callable, a: float, b: float) -> float:
- 89 # Estimate \int_{a}^{b} f^{2}(x) dx
- 90 E, _ = spi.quadrature(lambda _: f(_) ** 2, a, b, maxiter=200)
- 91 return E
- 92
- 93 if np.isclose(eps, 0):
- 94 sF = np.inf
- 95 else:
- 96 s = self.support()
- 97 E_tot = energy(self.apply, -s, s)
- 98
- 99 # Coarse-grain search for a max bandwidth in v_step increments.
-100 tolerance_reached = False
-101 v_step = 1 / s # slowest decay signal is sinc() -> steps at its zeros.
-102 v_max = 0
-103 while not tolerance_reached:
-104 v_max += v_step
-105 E = energy(self.applyF, -v_max, v_max)
-106 tolerance_reached = E >= tol * E_tot
-107
-108 # Fine-grained search for a max bandwidth in [v_max - v_step, v_max] region.
-109 v_fine = np.linspace(v_max - v_step, v_max, 100)
-110 E = np.array([energy(self.applyF, -v, v) for v in v_fine])
-111
-112 sF = v_fine[E >= tol * E_tot].min()
-113 return sF
-
-114
-115 def argscale(self, scalar: pxt.Real) -> pxt.OpT:
-116 scalar = float(scalar)
-117 assert scalar > 0
-118
-119 @pxrt.enforce_precision(i="arr")
-120 def g_apply(_, arr: pxt.NDArray) -> pxt.NDArray:
-121 # :math:`g(x) = f(\alpha x)`
-122 op, cst = _._op, _._cst
-123 return op.apply(cst * arr)
-124
-125 def g_support(_) -> pxt.Real:
-126 op, cst = _._op, _._cst
-127 return op.support() / cst
-128
-129 @pxrt.enforce_precision(i="arr")
-130 def g_applyF(_, arr: pxt.NDArray) -> pxt.NDArray:
-131 # :math:`g^{F}(v) = f^{F}(v / \alpha) / \alpha`
-132 op, cst = _._op, _._cst
-133 return op.applyF(arr / cst) / cst
-134
-135 def g_supportF(_, eps: pxt.Real) -> pxt.Real:
-136 op, cst = _._op, _._cst
-137 return op.supportF(eps) * cst
-138
-139 g = FSSPulse()
-140 g._cst = scalar # scale factor
-141 g._op = self # initial pulse
+[docs]
+ 67 def supportF(self, eps: pxt.Real) -> pxt.Real:
+ 68 r"""
+ 69 Parameters
+ 70 ----------
+ 71 eps: Real
+ 72 Energy cutoff threshold :math:`\epsilon \in [0, 0.05]`.
+ 73
+ 74 Returns
+ 75 -------
+ 76 sF: Real
+ 77 Value such that
+ 78
+ 79 .. math::
+ 80
+ 81 \int_{-s^{\mathcal{F}}}^{s^{\mathcal{F}}} |f^{\mathcal{F}}(v)|^{2} dv
+ 82 \approx
+ 83 (1 - \epsilon) \|f\|_{2}^{2}
+ 84 """
+ 85 eps = float(eps)
+ 86 assert 0 <= eps <= 0.05
+ 87 tol = 1 - eps
+ 88
+ 89 def energy(f: callable, a: float, b: float) -> float:
+ 90 # Estimate \int_{a}^{b} f^{2}(x) dx
+ 91 E, _ = spi.quadrature(lambda _: f(_) ** 2, a, b, maxiter=200)
+ 92 return E
+ 93
+ 94 if np.isclose(eps, 0):
+ 95 sF = np.inf
+ 96 else:
+ 97 s = self.support()
+ 98 E_tot = energy(self.apply, -s, s)
+ 99
+100 # Coarse-grain search for a max bandwidth in v_step increments.
+101 tolerance_reached = False
+102 v_step = 1 / s # slowest decay signal is sinc() -> steps at its zeros.
+103 v_max = 0
+104 while not tolerance_reached:
+105 v_max += v_step
+106 E = energy(self.applyF, -v_max, v_max)
+107 tolerance_reached = E >= tol * E_tot
+108
+109 # Fine-grained search for a max bandwidth in [v_max - v_step, v_max] region.
+110 v_fine = np.linspace(v_max - v_step, v_max, 100)
+111 E = np.array([energy(self.applyF, -v, v) for v in v_fine])
+112
+113 sF = v_fine[E >= tol * E_tot].min()
+114 return sF
+
+115
+116 def argscale(self, scalar: pxt.Real) -> pxt.OpT:
+117 scalar = float(scalar)
+118 assert scalar > 0
+119
+120 @pxrt.enforce_precision(i="arr")
+121 def g_apply(_, arr: pxt.NDArray) -> pxt.NDArray:
+122 # :math:`g(x) = f(\alpha x)`
+123 op, cst = _._op, _._cst
+124 return op.apply(cst * arr)
+125
+126 def g_support(_) -> pxt.Real:
+127 op, cst = _._op, _._cst
+128 return op.support() / cst
+129
+130 @pxrt.enforce_precision(i="arr")
+131 def g_applyF(_, arr: pxt.NDArray) -> pxt.NDArray:
+132 # :math:`g^{F}(v) = f^{F}(v / \alpha) / \alpha`
+133 op, cst = _._op, _._cst
+134 return op.applyF(arr / cst) / cst
+135
+136 def g_supportF(_, eps: pxt.Real) -> pxt.Real:
+137 op, cst = _._op, _._cst
+138 return op.supportF(eps) * cst
+139
+140 def g_expr(_) -> tuple:
+141 return ("argscale", _._op, _._cst)
142
-143 g.apply = types.MethodType(g_apply, g)
-144 g.support = types.MethodType(g_support, g)
-145 g.applyF = types.MethodType(g_applyF, g)
-146 g.supportF = types.MethodType(g_supportF, g)
-147 return g