diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 8a0fb9a..0956746 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -52,7 +52,7 @@ jobs: galsim \ ngmix - python -m pip install tensorflow==2.5.0 tensorflow_addons==0.13.0 tensorflow_probability==0.13.0 + python -m pip install tensorflow==2.5.0 tensorflow_addons==0.13.0 tensorflow_probability==0.13.0 numdifftools python -m pip install git+https://github.com/DifferentiableUniverseInitiative/GalFlow.git python -m pip install -e . - name: test diff --git a/tests/test_interpolation_gradients.py b/tests/test_interpolation_gradients.py new file mode 100644 index 0000000..47670fd --- /dev/null +++ b/tests/test_interpolation_gradients.py @@ -0,0 +1,61 @@ +# This module tests tfa gradients in respect to interpolation methods. +import numpy as np +from tensorflow_addons.image import resampler +from scipy.misc import face +import numdifftools +import tensorflow as tf + +from numpy.testing import assert_allclose + + +def facer(interpolant, warp_tf): + image = face(gray=True)[-512:-512+128,-512:-512+128].astype('float32') + image_tf = tf.convert_to_tensor(image.reshape([1,128,128, 1])) + #define a shift + shift = tf.zeros([1,2]) + + #calculate derivatives via tf.GradientTape + with tf.GradientTape() as tape: + tape.watch(shift) + ws = tf.reshape(shift,[1,1,1,2]) + warp_tf + o = resampler(image_tf, ws, interpolant) + autodiff_jacobian = tape.batch_jacobian(o, shift) + + #calculate derivatives via numdifftools + def fn(shift): + shift = tf.convert_to_tensor(shift.astype('float32')) + ws = tf.reshape(shift,[1,1,1,2]) + warp_tf + o = resampler(image_tf, ws, interpolant) + return o.numpy().flatten() + + numdiff_jacobian = numdifftools.Jacobian(fn, order=4, step=0.04) + numdiff_jacobian = numdiff_jacobian(np.zeros([2])).reshape([128,128,2]) + + return autodiff_jacobian[0,...,0,:], numdiff_jacobian + +import pytest + +xfail = pytest.mark.xfail +@xfail(reason="Fails because it needs the modified tensorflow_addons to work") +def test_interpolation_gradients(): + atol = 0.003 #taken from the bilinear case with half step warp. + + interpolant = "bilinear" + #on pixel interpolation + int_warp = np.stack(np.meshgrid(np.arange(128), np.arange(128)), axis=-1).astype('float32') + int_warp_tf = tf.convert_to_tensor(int_warp.reshape([1,128,128,2])) + + #half step interpolation + half_warp = np.stack(np.meshgrid(np.arange(128), np.arange(128)), axis=-1).astype('float32') + half_warp_tf = tf.convert_to_tensor(half_warp.reshape([1,128,128,2])+.5) #add a half-step + + autodiff_jacobian_int, numdiff_jacobian_int = facer(interpolant,int_warp_tf) + autodiff_jacobian_half, numdiff_jacobian_half = facer(interpolant,half_warp_tf) + + + assert_allclose(autodiff_jacobian_half,numdiff_jacobian_half, rtol=0.1, atol=atol) + assert_allclose(autodiff_jacobian_int, numdiff_jacobian_int, rtol=0.1, atol=atol) + + +if __name__=='__main__': + test_interpolation_gradients()