Numerical error on division (np vs jnp)

[JLjw8]

Hi,

We are wondering if there is a way to minimize the division numerical error between np and jnp, when the denominator and nummerator are almost 0. We have tried to cast the jax.numpy to use float64 so as to match with numpy, but it still cannot handle the following line as expected:

numpy version:
cos_az = ((np.sin(latlng.lat().radians) * np.cos(zenith_angle) - np.sin(declination_sun)) / (np.cos(latlng.lat().radians) * np.sin(zenith_angle)))

jnp version:
cos_az = (jnp.sin(latlng.lat) * jnp.cos(zenith_angle) - jnp.sin(declination_sun)) / (jnp.cos(latlng.lat) * jnp.sin(zenith_angle))

The values being passed in to both lines are the same, but the results vary too much.

Input:

• np: latlng.lat().radians: -1.5707963267948966, zenith_angle: 1.161716459429606, declination_sun: -0.40907986736529034
• jnp: latlng.lat: -1.5707963267948966, zenith_angle: 1.1617164594296063, declination_sun: -0.4090798673652903

Intermediate computation in cos_az:

• np: (-1.6653345369377348e-16) / (5.617988862548204e-17)
• jnp: (0.0) / (5.617988862548204e-17)

Output:

• np: cos_az=-2.9642894951955694
• jnp: cos_az=0.0

[jakevdp]

Hi - thanks for the question! It would be much easier to diagnose what's happening and answer your question if you could create a minimal reproducible example: i.e a block of code that others could run to reproduce the results you're seeing. For example, I have no idea what latlng.lat is supposed to be. Thanks!

[JLjw8]

Hi! Please find our minimal reproducible example below. We found out that it was because of the error in arcsin that made more and more precision errors aggregate. The first test works, and the rest of the tests are either off by one digit, drop one digit, or off by two digits. Our jax version is 0.4.30, and our Python version is 3.9.20.

import numpy as np
import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)

def numpy_version(obliquity_correction, apparent_long_sun):
    declination_sun = np.arcsin(
      np.sin(obliquity_correction) * np.sin(apparent_long_sun))
    return declination_sun

def jnp_version(obliquity_correction, apparent_long_sun):
    declination_sun = jnp.arcsin(
      jnp.sin(obliquity_correction) * jnp.sin(apparent_long_sun))
    return declination_sun

tests = [ (0.409080651134202, 157.10346628729434), (0.4090810491912244, 160.237885151924), (0.4090807783600018, 161.79407168447165), (0.40908100951762, 158.66099390508788) ] 
functions = [numpy_version, jnp_version ]

for i, test in enumerate(tests):
    print(f'Test {i+1}:')
    print(f"\tinput= {test}")
    for func in functions:
        print(f'\t{func.__name__}={func(*test)}')

Our output is as follows:

Test 1:
	input = (0.409080651134202, 157.10346628729434)
	numpy_version=0.009479443448665083
	jnp_version=0.009479443448665083
Test 2:
	input = (0.4090810491912244, 160.237885151924)
	numpy_version=-0.006626457639579936
	jnp_version=-0.0066264576395799355
Test 3:
	input = (0.4090807783600018, 161.79407168447165)
	numpy_version=-0.40907986736529034
	jnp_version=-0.4090798673652903
Test 4:
	input = (0.40908100951762, 158.66099390508788)
	numpy_version=0.40905681468675686
	jnp_version=0.4090568146867569

[jakevdp]

Thanks - that's helpful. It looks like the results of jnp.arcsin differ from those of np.arcsin at about 1 part in $10^{15}$ or $10^{16}$, which is close to the unit in the last place. You should keep in mind that at this level of accuracy, numpy is not an entirely reliable baseline to be using for judging correctness.

I'm going to ping @pearu who has been working on the accuracy of arcsin in various domains, and may have some insight here.

[pearu]

First, the original example in the issue and the provided reproducer are not related with respect to the source of numerical errors. While JAX and numpy apparently use different algorithms for arcsin (or these use different FTZ modes), the differences between results are within acceptable limits (less than or equal to 1 ULP).

The example tries to evaluate cos_az close to (south?) pole where its expression is 0/0 indeterminate. It is not clear from the post which of the results is expected. So, when applying L’Hopital’s Rule for the limit process lat->-pi/2, I conclude that the expected result of cos_az approaches to 0. This matches with JAX results. On the other hand, I would not blame numpy either because the evaluation of an expression close to its singularity point using floating-point numbers most likely returns garbage anyway.

I recommend splitting the evaluation of cos_az into two pieces. When lat is far from +-pi/2, then use the original expression. And when lat is close to poles (say, when abs(pi/2 - abs(lat)) < 1e-5), use a limiting expression (-0.5 * cos(lat) * cos(zenith_angle) / (sin(lat) * sin(zenith_angle)) seems to give correct results).

[JLjw8]

Thank you!!