DifferentiableUniverseInitiative/jax_cosmo

[TypeError] Ensure children is a list before calling children.pop() for a cosmology with the growth index $\gamma$

Opened this issue · 0 comments

Hi,

I was looking into jax-cosmo for a fisher forecast of future constraints on the growth index $\gamma$.

When I include $\gamma$ into the cosmological parameters in a jax_cosmo.Cosmology object, and try to compute the 3x2-point signal mean or covariance, I get the following error when the function jax_cosmo.angular_cl.angular_cl(cosmo, ell, probes) is called:

File /usr/local/anaconda/envs/py311forge/lib/python3.11/site-packages/jax_cosmo-0.1.dev262+g816069f-py3.11.egg/jax_cosmo/core.py:128, in Cosmology.tree_unflatten(cls, aux_data, children)
    125 # We extract the remaining parameters in reverse order from how they
    126 # were inserted
    127 if aux_data["gamma_growth"]:
--> 128     gamma = children.pop()
    129 else:
    130     gamma = None

AttributeError: 'NoneType' object has no attribute 'pop'

Looking at the source code in core.py, my guess is that the issue originates around these lines

 @classmethod
    def tree_unflatten(cls, aux_data, children):
        # Retrieve base parameters
        Omega_c, Omega_b, h, n_s, sigma8, Omega_k, w0, wa = children[:8]
        children = list(children[8:]).reverse()

        # We extract the remaining parameters in reverse order from how they
        # were inserted
        if aux_data["gamma_growth"]:
            gamma = children.pop()
        else:
            gamma = None

I think reverse() does not return a reversed list but rather reverses the list in-place, then returns None. I believe that leads to the error message I got. If that is indeed the case, an easy fix would be something along the line of

Omega_c, Omega_b, h, n_s, sigma8, Omega_k, w0, wa = children[:8]
    children = list(children[8:])
    children.reverse()

Do you want me to open a pull request?