LSSTDESC / lsstdesc-diffsky

Library for differentiable generation of synthetic skies for LSST DESC
BSD 3-Clause "New" or "Revised" License
1 stars 0 forks source link

Resolve deprecation warning in pecZ.py module #75

Closed aphearin closed 8 months ago

aphearin commented 8 months ago

This PR resolves a deprecation warning in pecZ.py by replacing the use of numpy.core.umath_tests.inner1d with a vectorized call to jnp.dot.

Previously, the pecZ function computed the peculiar velocity v_pec like this:

v_pec = numpy.core.umath_tests.inner1d(v, r_rel_hat)

The above line of code computes the elementwise dot product between v and r_rel_hat. Presumably, inner1d was used because np.dot only computes the dot product between a single pair of vectors, rather than an array of vectors. But importing from numpy.core.umath_tests triggers a deprecation warning notifying us that this will break in future. The previous implementation addressed this warning by ignoring it:

warnings.filterwarnings("ignore", category=DeprecationWarning)

This PR instead resolves this deprecation warning by computing the elementwise peculiar velocity by using JAX to vectorize jnp.dot as follows:

from jax import vmap, jit as jjit, numpy as jnp
dot_vmap = jjit(vmap(jnp.dot, in_axes=(0, 0)))
v_pec = dot_vmap(v, r_rel_hat)

This change was not straightforward to implement because pecZ.py had no testing. But there is now a new test_pecZ.py module that implements a basic unit test ensuring that at least the pecZ function returns finite-valued arrays of the expected shape.

CC @evevkovacs

codecov[bot] commented 8 months ago

Codecov Report

All modified lines are covered by tests :white_check_mark:

Comparison is base (85a1832) 61.26% compared to head (f7da55a) 61.48%.

Additional details and impacted files ```diff @@ Coverage Diff @@ ## main #75 +/- ## ========================================== + Coverage 61.26% 61.48% +0.21% ========================================== Files 86 86 Lines 5298 5299 +1 ========================================== + Hits 3246 3258 +12 + Misses 2052 2041 -11 ``` | [Files](https://app.codecov.io/gh/LSSTDESC/lsstdesc-diffsky/pull/75?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=LSSTDESC) | Coverage Δ | | |---|---|---| | [lsstdesc\_diffsky/pecZ.py](https://app.codecov.io/gh/LSSTDESC/lsstdesc-diffsky/pull/75?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=LSSTDESC#diff-bHNzdGRlc2NfZGlmZnNreS9wZWNaLnB5) | `83.33% <100.00%> (+48.55%)` | :arrow_up: |

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.