Spergel

[jecampagne]

This is an implementation of the Spergel profile.
In the Galsim code the Spergel index (nu) is restricted to [-0.85, 4.0]. In the JAX code I do not know to assert a use of non valid nu input parameter.

Notice that nu=0.5 is exactly the Exponential profile. Spergel profile might be a good alternative to the Sersic profile as realized here.

The present code pass test_spergel.py Galsim tests excepting the photon shooting even if I have setup a first version using a inversion method using JAX-Galsim bisection root algo.

A side remark. For Moffat & Spergel we need K_nu Bessel function. For the time beeing this Bessel func is defined both in moffat.py and spergel.py. I guess this is not optimal and this function would have to be defined in the core directory.

[jecampagne]

I do not understand such message in the pytest used in the PR causing test failure

" test_api_gsobject[docs-methods]: 
self = <[AttributeError("'Spergel' object has no attribute 'beta'") raised in repr()] Spergel object at 0x7f77ae66dd90>

[beckermr]

I found a few small typos that may help get the tests to pass.

[jecampagne]

Now test_api.py have detected that my xValue is not/ jit-grad friendly. I must investigate.

After some tests the problem seems to be in fsmallz_nu(z, nu) as

def new_xValue(self, pos):
        r = jnp.sqrt(pos.x**2 + pos.y**2) * self._inv_r0
        #jnp.where(z <= 1.0e-10, fsmallz_nu(z, nu), jnp.power(z, nu) * _Knu(nu, z))
        res = jnp.where(r <= 1.0e-10, 
                        jnp.power(r, self.nu) * galsim.spergel._Knu(self.nu, r),
                        jnp.power(r, self.nu) * galsim.spergel._Knu(self.nu, r)
        )
        return self._xnorm * res
galsim.spergel.Spergel._xValue = new_xValue

Is ok.

[jecampagne]

Well I'm a bit puzzled but I guess this is related to
https://jax.readthedocs.io/en/latest/faq.html#gradients-contain-nan-where-using-where

in fact fsmallz_nu cannot be used with integer nu so then the gardient. So I need to protect it according to the above doc.
The pb is that jnp.power(z, nu) * _Knu(nu, z) is not accurate for small value of z which is a pb for very sharp Spergel profile.

A working hack is (here a code tested in Google colab)

def new_fsmallz_nu(z, nu):
  def fnu(z, nu):
        """z^nu K_nu[z] z -> 0 O(z^4) z > 0"""
        nu += 1.e-10 #to garanty that nu is not an integer
        z2 = z * z
        z4 = z2 * z2
        c1 = jnp.power(2.0, -6.0 - nu)
        c2 = galsim.spergel._gamma(-2.0 - nu)
        c3 = galsim.spergel._gamma(-2.0 + nu)
        c4 = jnp.power(z, 2.0 * nu)
        c5 = z4 * 8.0 * z2 * (2.0 + nu) + 32.0 * (1.0 + nu) * (2.0 + nu)
        c6 = z2 * (16.0 + z2 - 8.0 * nu) * c3
        return c1 * (c4 * c5 * c2 + jnp.power(4.0, nu) * (c6 + 32.0 * galsim.spergel._gamma(nu)))
  return jnp.select(
        [nu == 0, nu == 1, nu == 2, nu == 3, nu == 4],
        [galsim.spergel.f0(z), galsim.spergel.f1(z), fgalsim.spergel.2(z), galsim.spergel.f3(z), galsim.spergel.f4(z)],
        default=fnu(z,nu),
  )

notice that in fact the NaN originate from the jnp.where used in

def fz_nu(z, nu):
    """z^nu K_nu[z], z > 0"""
    return jnp.where(z <= 1.0e-10, fsmallz_nu(z, nu), jnp.power(z, nu) * _Knu(nu, z))

[jecampagne]

I have made progress to get test_serpel.py::test_spergel do_shoot working.
Now in the test_serpel.py::test_spergel_shoot : nu >0 tests are ok. The test with

 obj = galsim.Spergel(nu=-0.6, half_light_radius=3.5, flux=1.e4)

fails as added_flux from

added_flux, photons = obj.drawPhot(im, poisson_flux=False, rng=rng.duplicate())

differs from obj.flux:

obj.flux =  10000.0
added_flux =  9999.0

Now, looking at SBSpergel.cpp::SpergelInfo::shoot C++ function Galsim comment is in the context of nu<=0

     // exact s.b. profile diverges at origin, so replace the inner most circle
     // (defined such that enclosed flux is shoot_acccuracy) with a linear function
    // that contains the same flux and has the right value at r = rmin.
    // So need to solve the following for a and b:
    // int(2 pi r (a + b r) dr, 0..rmin) = shoot_accuracy
    // a + b rmin = K_nu(rmin) * rmin^nu

My comment is the following:

1. The Spergel profile noted I_{nu}(r) = xnorm(nu) * f(r/r0,nu) is such that

$$\Large f(z,\nu)=z^\nu K_\nu(z)$$

as the following property for z->0 (or r->0 for I_{nu}(r))

$$ \begin{cases}\Large f(z,\nu>0) = 2^{\nu -1} \Gamma(\nu)\\\ \Large f(z,\nu=0) \propto \log(1/z) \\\ \Large f(z,\nu<0) \propto 1/z^{-2\nu} \end{cases} $$

So, I agree that I_{nu}(r) is diverging for nu<=0.
2) But the cumulated integrated flux is proportional to

$$\Large \int_0^z f(z,\nu) z dz$$

has no divergence at z=0 as nu>-1 (which is the original paper bound) as nu is in [-0.85, 4.0] (GalSim bounds).

So, I do not see the need to make the procedure mentionned in GalSim and I do suspect that it introduces the difference detected by pytest.

Now, despite what I've written above, I will try to implement the GalSim procedure.

[jecampagne]

I have opened an issue in the GalSim repo to understand how they manage the case nu<0 as I may found a pb in the normalisation. Let see what Mike will tell me.

[jecampagne]

I have coded shoot for nu<=0 "as it is" in Galsim. The code pass Galsim/tests/test_spergel.py (ie. pytest -k test_spergel.py`). Now for me there is a pb with this code (I do not speak about optimisation as mentioned by Matthew):

1. There is no reason to do not use a inversion method of the cumulative flux for nu<=0 as for nu>=-0.85 (Galsim bound) the

$$\int_ 0^r u I(u) du$$

is a non-diverging integral when r->0.

2. The Galsim way to treat the generation r when nu<=0 relies on a linear approximation of I(r) = a + b r when r<rmin with rmin defined by the shoot_accuracy parameter. Going throw the math and looking the Galsim implementation I have detected a normalization error that leads to wrong (a,b) numerical determination. The error leads to the 0.1% error on the flux detected by pytest.

Now either we stick to the Galsim coding because I may be wrong in interpreting the Galsim coding, or we adopt a more natural way to proceed as from point (1) there is no matter to distinguish nu>0 and nu<=0.
After deciding, I have to optimize the shoot process along Matthiew remark. (Rq. the full pytest takes 40min or so).

[jecampagne]

Hi, I have implemented a shoot_neg(ie. for nu<0) such that 1) it follows the Galsim spirit of using a linear approx of the profile for low r value and 2) it uses a normalization which passes the pytest and is "math. correct".
I have also use the inversion of the cumulative flux along the exponential use-case (ie. lazy_property deco & linear interpolation inversion). The whole pytest last 23min now.

[jecampagne]

There is still a missing test to assert that nu parameter is in [-0.86, 4.0].

[beckermr]

I'll take a look tomorrow!

[beckermr]

I left one suggestion for rephrasing the docs.

My final comment is that the galsim APIs provide some bessel functions in galsim.bessel. Any bessel functions defined in this file that are provided by galsim in the galsim.bessel module should be moved there.

[jecampagne]

I have rename "Knu" (modified Bessel 2nd kind) as "kv" to follow Galsim notation. "kv" is moved to jax_galsim/bessel.py. I have realised then that Knu/kv was defined & used for Moffat, so I have remove the code to use the new one. I have not moved other z^nu Knu(nu,z) related functions as there are specific to Spergel profile, nor the Gamma function.

Now in the coredirectory there is a "bessel.py" file that I had created in the past. I think this should be removed and the code moved to jax_galsim/bessel.py. What do you think?

[beckermr]

Now in the coredirectory there is a "bessel.py" file that I had created in the past. I think this should be removed and the code moved to jax_galsim/bessel.py. What do you think?

No we should not do this. The goal is to match the galsim API exactly. So we only put bessel functions that galsim provides in jax_galsim/bessel.py. Anything extra, we can add to core since this contains private APIs to jax_galsim.

[jecampagne]

Who press the Merge pull request?

[beckermr]

Either! I'll go ahead!