Rational of 2 polynomials with gradient

[jecampagne]

Hello

For some personal use I had the compute P(x)/Q(x). I have found Boost C++ code that could be nicely be implemented in JAX. Here is my code, certainly not the best implementation

  def _evaluate_rational(z, num, denom):

    count = len(num)

    def true_fn_update(z):

      def body_true(val1):
        # decode val1
        s1, s2, i = val1 
        s1 *= z; s2 *= z; s1 += num[i]; s2 += denom[i];
        return s1, s2, i-1

      def cond_true(val1):
        s1, s2, i = val1 
        return i>=0

      val1_init = (num[-1], denom[-1], count-2)
      s1, s2, _ = jax.lax.while_loop(cond_true, body_true, val1_init)

      return s1/s2

    def false_fn_update(z):
      def body_false(val1):
        # decode val1
        s1, s2, i = val1 
        s1 *= z; s2 *= z; s1 += num[i]; s2 += denom[i];
        return s1, s2, i+1

      def cond_false(val1):
        s1, s2, i = val1 
        return i<count

      val1_init = (num[0], denom[0],1)
      s1, s2, _ = jax.lax.while_loop(cond_false, body_false, val1_init)

      return s1/s2


    return jnp.where(z<=1, true_fn_update(z), false_fn_update(1/z))  

Some tests

num = jnp.array([1.,2.,4.,-5.])
denom = jnp.array([-0.5,2.,-3,-10.])


def Ratio(x):
  return (1. + 2*x + 4*x*x - 5*x*x*x)/(-0.5 + 2*x -3.*x*x - 10.*x*x*x)

v_ratio = jax.vmap(_evaluate_rational, in_axes=(0,None, None))

x=jnp.linspace(0,100,10000)

plot(x,Ratio(x))
plot(x,v_ratio(x, num,denom),ls="--")
xscale("log")

leads to

For the gradiant

vgrad_ratio = jax.vmap(jax.jacfwd(_evaluate_rational), in_axes=(0,None, None))

def GradR(x):
  return (2 + 8*x - 15*x**2)/(-0.5 + 2*x - 3.*x**2 - 10.*x**3) - \
     -  ((2 - 6.*x - 30.*x**2)*(1. + 2*x + 4*x**2 - 5*x**3))/ \
     -   (-0.5 + 2*x - 3.*x**2 - 10.*x**3)**2

plot(x,GradR(x))
plot(x,vgrad_ratio(x, num,denom),ls="--")
xscale("log")

leads to

Now, I guess my code is not 100% as good as JAX.lax standard function inmplementation with the shape promotion, dtype.... But looks pretty a version 0, isn't it?

[jecampagne]

Does this code can be submitted to JAX improvement? If so what is the procedure as I do not see where in the JAX directory there is a place to make a PR ? I have also a Clenshaw-Curtis integration quadrature routine if desired...

[jecampagne]

Hi,
In fact it is not exactlty the same thing
if jnp.polydiv(num,den) provide q,r then num(x)/den(x) = q(x) + r(x)/den(x).

But, I agree that I was producing an overkill code !

def R(z, num,denom):
  return jnp.polyval(num,z)/jnp.polyval(denom,z)

is perfectly good a solution. Now, what a bout the speed efficicency I do not know.. If Boost as used a specific implementation of the ratio, I guess that the Horner computation of the numerator & denominator has less efficiency than computing the ratio as I have done in JAX...

[jakevdp]

OK I see... I think unless the algorithm you implemented is something provided by numpy or scipy, then it is probably out of scope for inclusion in jax itself. That said, I don't know of any good alternate package for code like this.

[jecampagne]

I agree that the "R" function above may be enough and do not need to be in the lib.

[jecampagne]

I add a small test. Let consider

num = jnp.array([1.,2.,4.,-5.])
denom = jnp.array([-0.5,2.,-3,-10.])

x=jnp.linspace(0,100,10000)

# here is a direct ratio computation
@jax.jit
def Ratio(x):
  return (1. + 2*x + 4*x*x - 5*x*x*x)/(-0.5 + 2*x -3.*x*x - 10.*x*x*x)

# here is a computiation of the ratio with call of jnp.polyval for both numerator and denominator
@jax.jit
def R(z, num,denom):
  return jnp.polyval(num,z)/jnp.polyval(denom,z)

# here is the JIT/vmap version of my implementation
v_ratio = jax.jit(jax.vmap(_evaluate_rational, in_axes=(0,None, None)))

After warmup phase to avoid compilation time

%timeit Ratio(x).block_until_ready()

gives 51.6 µs ± 12.2 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit R(x,num[::-1],denom[::-1]).block_until_ready()

gives 561 µs ± 67.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

and

%timeit v_ratio(x, num,denom).block_until_ready()

gives 78.9 µs ± 8.2 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

which tell us that the Horner's scheme applied on both polynomials is slower that the Boost algo. Of course the direct method is faster here as I guess there is optimizations that cannot be done in the general case.

Numpy does not seems to implement an optimized code for ratio of polynomials calculus.

[jakevdp]

I see, thanks. This does look like useful code, but I think it’s out of scope for the Jax package itself.