Coverage for src/bartz/jaxext/scipy/special.py: 95%
60 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-27 14:46 +0000
1# bartz/src/bartz/jaxext/scipy/special.py
2#
3# Copyright (c) 2025, Giacomo Petrillo
4#
5# This file is part of bartz.
6#
7# Permission is hereby granted, free of charge, to any person obtaining a copy
8# of this software and associated documentation files (the "Software"), to deal
9# in the Software without restriction, including without limitation the rights
10# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
11# copies of the Software, and to permit persons to whom the Software is
12# furnished to do so, subject to the following conditions:
13#
14# The above copyright notice and this permission notice shall be included in all
15# copies or substantial portions of the Software.
16#
17# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
20# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
22# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
23# SOFTWARE.
25"""Mockup of the :external:py:mod:`scipy.special` module."""
27from functools import wraps 1ab
29from jax import ShapeDtypeStruct, pure_callback 1ab
30from jax import numpy as jnp 1ab
31from scipy.special import gammainccinv as scipy_gammainccinv 1ab
34def _float_type(*args): 1ab
35 """Determine the jax floating point result type given operands/types."""
36 t = jnp.result_type(*args) 1ab
37 return jnp.sin(jnp.empty(0, t)).dtype 1ab
40def _castto(func, dtype): 1ab
41 @wraps(func) 1ab
42 def newfunc(*args, **kw): 1ab
43 return func(*args, **kw).astype(dtype) 1ab
45 return newfunc 1ab
48def gammainccinv(a, y): 1ab
49 """Survival function inverse of the Gamma(a, 1) distribution."""
50 a = jnp.asarray(a) 1ab
51 y = jnp.asarray(y) 1ab
52 shape = jnp.broadcast_shapes(a.shape, y.shape) 1ab
53 dtype = _float_type(a.dtype, y.dtype) 1ab
54 dummy = ShapeDtypeStruct(shape, dtype) 1ab
55 ufunc = _castto(scipy_gammainccinv, dtype) 1ab
56 return pure_callback(ufunc, dummy, a, y, vmap_method='expand_dims') 1ab
59################# COPIED AND ADAPTED FROM JAX ##################
60# Copyright 2018 The JAX Authors.
61#
62# Licensed under the Apache License, Version 2.0 (the "License");
63# you may not use this file except in compliance with the License.
64# You may obtain a copy of the License at
65#
66# https://www.apache.org/licenses/LICENSE-2.0
67#
68# Unless required by applicable law or agreed to in writing, software
69# distributed under the License is distributed on an "AS IS" BASIS,
70# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
71# See the License for the specific language governing permissions and
72# limitations under the License.
74import numpy as np 1ab
75from jax import debug_infs, lax 1ab
78def ndtri(p): 1ab
79 """Compute the inverse of the CDF of the Normal distribution function.
81 This is a patch of `jax.scipy.special.ndtri`.
82 """
83 dtype = lax.dtype(p) 1ab
84 if dtype not in (jnp.float32, jnp.float64): 84 ↛ 85line 84 didn't jump to line 85 because the condition on line 84 was never true1ab
85 msg = f'x.dtype={dtype} is not supported, see docstring for supported types.'
86 raise TypeError(msg)
87 return _ndtri(p) 1ab
90def _ndtri(p): 1ab
91 # Constants used in piece-wise rational approximations. Taken from the cephes
92 # library:
93 # https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
94 p0 = list( 1ab
95 reversed(
96 [
97 -5.99633501014107895267e1,
98 9.80010754185999661536e1,
99 -5.66762857469070293439e1,
100 1.39312609387279679503e1,
101 -1.23916583867381258016e0,
102 ]
103 )
104 )
105 q0 = list( 1ab
106 reversed(
107 [
108 1.0,
109 1.95448858338141759834e0,
110 4.67627912898881538453e0,
111 8.63602421390890590575e1,
112 -2.25462687854119370527e2,
113 2.00260212380060660359e2,
114 -8.20372256168333339912e1,
115 1.59056225126211695515e1,
116 -1.18331621121330003142e0,
117 ]
118 )
119 )
120 p1 = list( 1ab
121 reversed(
122 [
123 4.05544892305962419923e0,
124 3.15251094599893866154e1,
125 5.71628192246421288162e1,
126 4.40805073893200834700e1,
127 1.46849561928858024014e1,
128 2.18663306850790267539e0,
129 -1.40256079171354495875e-1,
130 -3.50424626827848203418e-2,
131 -8.57456785154685413611e-4,
132 ]
133 )
134 )
135 q1 = list( 1ab
136 reversed(
137 [
138 1.0,
139 1.57799883256466749731e1,
140 4.53907635128879210584e1,
141 4.13172038254672030440e1,
142 1.50425385692907503408e1,
143 2.50464946208309415979e0,
144 -1.42182922854787788574e-1,
145 -3.80806407691578277194e-2,
146 -9.33259480895457427372e-4,
147 ]
148 )
149 )
150 p2 = list( 1ab
151 reversed(
152 [
153 3.23774891776946035970e0,
154 6.91522889068984211695e0,
155 3.93881025292474443415e0,
156 1.33303460815807542389e0,
157 2.01485389549179081538e-1,
158 1.23716634817820021358e-2,
159 3.01581553508235416007e-4,
160 2.65806974686737550832e-6,
161 6.23974539184983293730e-9,
162 ]
163 )
164 )
165 q2 = list( 1ab
166 reversed(
167 [
168 1.0,
169 6.02427039364742014255e0,
170 3.67983563856160859403e0,
171 1.37702099489081330271e0,
172 2.16236993594496635890e-1,
173 1.34204006088543189037e-2,
174 3.28014464682127739104e-4,
175 2.89247864745380683936e-6,
176 6.79019408009981274425e-9,
177 ]
178 )
179 )
181 dtype = lax.dtype(p).type 1ab
182 shape = jnp.shape(p) 1ab
184 def _create_polynomial(var, coeffs): 1ab
185 """Compute n_th order polynomial via Horner's method."""
186 coeffs = np.array(coeffs, dtype) 1ab
187 if not coeffs.size: 1ab
188 return jnp.zeros_like(var) 1ab
189 return coeffs[0] + _create_polynomial(var, coeffs[1:]) * var 1ab
191 maybe_complement_p = jnp.where(p > dtype(-np.expm1(-2.0)), dtype(1.0) - p, p) 1ab
192 # Write in an arbitrary value in place of 0 for p since 0 will cause NaNs
193 # later on. The result from the computation when p == 0 is not used so any
194 # number that doesn't result in NaNs is fine.
195 sanitized_mcp = jnp.where( 1ab
196 maybe_complement_p == dtype(0.0),
197 jnp.full(shape, dtype(0.5)),
198 maybe_complement_p,
199 )
201 # Compute x for p > exp(-2): x/sqrt(2pi) = w + w**3 P0(w**2)/Q0(w**2).
202 w = sanitized_mcp - dtype(0.5) 1ab
203 ww = lax.square(w) 1ab
204 x_for_big_p = w + w * ww * (_create_polynomial(ww, p0) / _create_polynomial(ww, q0)) 1ab
205 x_for_big_p *= -dtype(np.sqrt(2.0 * np.pi)) 1ab
207 # Compute x for p <= exp(-2): x = z - log(z)/z - (1/z) P(1/z) / Q(1/z),
208 # where z = sqrt(-2. * log(p)), and P/Q are chosen between two different
209 # arrays based on whether p < exp(-32).
210 z = lax.sqrt(dtype(-2.0) * lax.log(sanitized_mcp)) 1ab
211 first_term = z - lax.log(z) / z 1ab
212 second_term_small_p = ( 1ab
213 _create_polynomial(dtype(1.0) / z, p2)
214 / _create_polynomial(dtype(1.0) / z, q2)
215 / z
216 )
217 second_term_otherwise = ( 1ab
218 _create_polynomial(dtype(1.0) / z, p1)
219 / _create_polynomial(dtype(1.0) / z, q1)
220 / z
221 )
222 x_for_small_p = first_term - second_term_small_p 1ab
223 x_otherwise = first_term - second_term_otherwise 1ab
225 x = jnp.where( 1ab
226 sanitized_mcp > dtype(np.exp(-2.0)),
227 x_for_big_p,
228 jnp.where(z >= dtype(8.0), x_for_small_p, x_otherwise),
229 )
231 x = jnp.where(p > dtype(1.0 - np.exp(-2.0)), x, -x) 1ab
232 with debug_infs(False): 1ab
233 infinity = jnp.full(shape, dtype(np.inf)) 1ab
234 neg_infinity = -infinity 1ab
235 return jnp.where( 1ab
236 p == dtype(0.0), neg_infinity, jnp.where(p == dtype(1.0), infinity, x)
237 )
240################################################################