Remove __future__ imports

Since Python2 is not supported by tis program the `__future__` imports are irrelevant.
LuzGarciaP · Sep 28, 2020 · 60e51f2 · 60e51f2
1 parent 842225f
commit 60e51f2
Show file tree

Hide file tree

Showing 14 changed files with 32 additions and 98 deletions.
diff --git a/jax_cosmo/__init__.py b/jax_cosmo/__init__.py
@@ -1,8 +1,4 @@
 # Cosmology in JAX
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 from pkg_resources import DistributionNotFound
 from pkg_resources import get_distribution
 

diff --git a/jax_cosmo/angular_cl.py b/jax_cosmo/angular_cl.py
@@ -1,8 +1,4 @@
 # This module contains functions to compute angular cls for various tracers
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 from functools import partial
 
 import jax.numpy as np

diff --git a/jax_cosmo/background.py b/jax_cosmo/background.py
@@ -1,8 +1,4 @@
 # This module implements various functions for the background COSMOLOGY
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import jax.numpy as np
 
 import jax_cosmo.constants as const
@@ -249,7 +245,7 @@ def dchioverdlna(y, x):
 
 
 def a_of_chi(cosmo, chi):
-    r""" Computes the scale factor for corresponding (array) of radial comoving
+    r"""Computes the scale factor for corresponding (array) of radial comoving
     distance by reverse linear interpolation.
 
     Parameters:
@@ -366,7 +362,7 @@ def angular_diameter_distance(cosmo, a):
 
 
 def growth_factor(cosmo, a):
-    """ Compute linear growth factor D(a) at a given scale factor,
+    """Compute linear growth factor D(a) at a given scale factor,
     normalized such that D(a=1) = 1.
 
     Parameters
@@ -396,7 +392,7 @@ def growth_factor(cosmo, a):
 
 
 def growth_rate(cosmo, a):
-    """ Compute growth rate dD/dlna at a given scale factor.
+    """Compute growth rate dD/dlna at a given scale factor.
 
     Parameters
     ----------
@@ -438,7 +434,7 @@ def growth_rate(cosmo, a):
 
 
 def _growth_factor_ODE(cosmo, a, log10_amin=-3, steps=128, eps=1e-4):
-    """ Compute linear growth factor D(a) at a given scale factor,
+    """Compute linear growth factor D(a) at a given scale factor,
     normalised such that D(a=1) = 1.
 
     Parameters
@@ -486,7 +482,7 @@ def D_derivs(y, x):
 
 
 def _growth_rate_ODE(cosmo, a):
-    """ Compute growth rate dD/dlna at a given scale factor by solving the linear
+    """Compute growth rate dD/dlna at a given scale factor by solving the linear
     growth ODE.
 
     Parameters
@@ -510,7 +506,7 @@ def _growth_rate_ODE(cosmo, a):
 
 
 def _growth_factor_gamma(cosmo, a, log10_amin=-3, steps=128):
-    r""" Computes growth factor by integrating the growth rate provided by the
+    r"""Computes growth factor by integrating the growth rate provided by the
     \gamma parametrization. Normalized such that D( a=1) =1
 
     Parameters

diff --git a/jax_cosmo/bias.py b/jax_cosmo/bias.py
@@ -1,8 +1,4 @@
 # This module contains implementations of galaxy bias
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import jax.numpy as np
 from jax.tree_util import register_pytree_node_class
 

diff --git a/jax_cosmo/core.py b/jax_cosmo/core.py
@@ -1,7 +1,3 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import jax.numpy as np
 from jax.experimental.ode import odeint
 from jax.tree_util import register_pytree_node_class

diff --git a/jax_cosmo/jax_utils.py b/jax_cosmo/jax_utils.py
@@ -1,7 +1,3 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import jax.numpy as np
 
 

diff --git a/jax_cosmo/likelihood.py b/jax_cosmo/likelihood.py
@@ -1,8 +1,4 @@
 # This module implements a few likelihoods useful
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import jax.numpy as np
 import jax.scipy as sp
 

diff --git a/jax_cosmo/parameters.py b/jax_cosmo/parameters.py
@@ -1,8 +1,4 @@
 # This module defines a few default cosmologies
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 from functools import partial
 
 from jax_cosmo.core import Cosmology

diff --git a/jax_cosmo/power.py b/jax_cosmo/power.py
@@ -1,8 +1,4 @@
 # This module computes power spectra
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import jax
 import jax.numpy as np
 
@@ -17,14 +13,14 @@
 
 
 def primordial_matter_power(cosmo, k):
-    """ Primordial power spectrum
-        Pk = k^n
+    """Primordial power spectrum
+    Pk = k^n
     """
     return k ** cosmo.n_s
 
 
 def linear_matter_power(cosmo, k, a=1.0, transfer_fn=tklib.Eisenstein_Hu, **kwargs):
-    r""" Computes the linear matter power spectrum.
+    r"""Computes the linear matter power spectrum.
 
     Parameters
     ----------
@@ -59,7 +55,7 @@ def linear_matter_power(cosmo, k, a=1.0, transfer_fn=tklib.Eisenstein_Hu, **kwar
 
 
 def sigmasqr(cosmo, R, transfer_fn, kmin=0.0001, kmax=1000.0, ksteps=5, **kwargs):
-    """ Computes the energy of the fluctuations within a sphere of R h^{-1} Mpc
+    """Computes the energy of the fluctuations within a sphere of R h^{-1} Mpc
 
     .. math::
 
@@ -84,15 +80,14 @@ def int_sigma(logk):
 
 
 def linear(cosmo, k, a, transfer_fn):
-    """Linear matter power spectrum
-    """
+    """Linear matter power spectrum"""
     return linear_matter_power(cosmo, k, a, transfer_fn)
 
 
 def _halofit_parameters(cosmo, a, transfer_fn):
-    r""" Computes the non linear scale,
-         effective spectral index,
-         spectral curvature
+    r"""Computes the non linear scale,
+    effective spectral index,
+    spectral curvature
     """
     # Step 1: Finding the non linear scale for which sigma(R)=1
     # That's our search range for the non linear scale
@@ -144,7 +139,7 @@ def integrand(logk):
 
 
 def halofit(cosmo, k, a, transfer_fn, prescription="takahashi2012"):
-    r""" Computes the non linear halofit correction to the matter power spectrum.
+    r"""Computes the non linear halofit correction to the matter power spectrum.
 
     Parameters
     ----------
@@ -271,7 +266,7 @@ def halofit(cosmo, k, a, transfer_fn, prescription="takahashi2012"):
 def nonlinear_matter_power(
     cosmo, k, a=1.0, transfer_fn=tklib.Eisenstein_Hu, nonlinear_fn=halofit
 ):
-    """ Computes the non-linear matter power spectrum.
+    """Computes the non-linear matter power spectrum.
 
     This function is just a wrapper over several nonlinear power spectra.
     """

diff --git a/jax_cosmo/probes.py b/jax_cosmo/probes.py
@@ -1,8 +1,4 @@
 # This module defines kernel functions for various tracers
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import jax.numpy as np
 from jax import jit
 from jax import vmap
@@ -187,7 +183,7 @@ def noise(self):
 
 @register_pytree_node_class
 class NumberCounts(container):
-    """ Class representing a galaxy clustering probe, with a bunch of bins
+    """Class representing a galaxy clustering probe, with a bunch of bins
 
     Parameters:
     -----------
@@ -214,14 +210,13 @@ def zmax(self):
 
     @property
     def n_tracers(self):
-        """ Returns the number of tracers for this probe, i.e. redshift bins
-        """
+        """Returns the number of tracers for this probe, i.e. redshift bins"""
         # Extract parameters
         pzs = self.params[0]
         return len(pzs)
 
     def kernel(self, cosmo, z, ell):
-        """ Compute the radial kernel for all nz bins in this probe.
+        """Compute the radial kernel for all nz bins in this probe.
 
         Returns:
         --------
@@ -235,7 +230,7 @@ def kernel(self, cosmo, z, ell):
         return kernel
 
     def noise(self):
-        """ Returns the noise power for all redshifts
+        """Returns the noise power for all redshifts
         return: shape [nbins]
         """
         # Extract parameters

diff --git a/jax_cosmo/redshift.py b/jax_cosmo/redshift.py
@@ -1,8 +1,4 @@
 # Module to define redshift distributions we can differentiate through
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 from abc import ABC
 from abc import abstractmethod
 
@@ -19,21 +15,18 @@
 
 class redshift_distribution(container):
     def __init__(self, *args, gals_per_arcmin2=1.0, zmax=10.0, **kwargs):
-        """ Initialize the parameters of the redshift distribution
-        """
+        """Initialize the parameters of the redshift distribution"""
         self._norm = None
         self._gals_per_arcmin2 = gals_per_arcmin2
         super(redshift_distribution, self).__init__(*args, zmax=zmax, **kwargs)
 
     @abstractmethod
     def pz_fn(self, z):
-        """ Un-normalized n(z) function provided by sub classes
-        """
+        """Un-normalized n(z) function provided by sub classes"""
         pass
 
     def __call__(self, z):
-        """ Computes the normalized n(z)
-        """
+        """Computes the normalized n(z)"""
         if self._norm is None:
             self._norm = simps(lambda t: self.pz_fn(t), 0.0, self.config["zmax"], 256)
         return self.pz_fn(z) / self._norm
@@ -44,15 +37,14 @@ def zmax(self):
 
     @property
     def gals_per_arcmin2(self):
-        """ Returns the number density of galaxies in gals/sq arcmin
+        """Returns the number density of galaxies in gals/sq arcmin
         TODO: find a better name
         """
         return self._gals_per_arcmin2
 
     @property
     def gals_per_steradian(self):
-        """ Returns the number density of galaxies in steradian
-        """
+        """Returns the number density of galaxies in steradian"""
         return self._gals_per_arcmin2 * steradian_to_arcmin2
 
     # Operations for flattening/unflattening representation
@@ -69,7 +61,7 @@ def tree_unflatten(cls, aux_data, children):
 
 @register_pytree_node_class
 class smail_nz(redshift_distribution):
-    """ Defines a smail distribution with these arguments
+    """Defines a smail distribution with these arguments
     Parameters:
     -----------
     a:
@@ -88,7 +80,7 @@ def pz_fn(self, z):
 
 @register_pytree_node_class
 class kde_nz(redshift_distribution):
-    """ A redshift distribution based on a KDE estimate of the nz of a
+    """A redshift distribution based on a KDE estimate of the nz of a
     given catalog currently uses a Gaussian kernel.
     TODO: add more if necessary
 
@@ -106,8 +98,7 @@ class kde_nz(redshift_distribution):
     """
 
     def _kernel(self, bw, X, x):
-        """ Gaussian kernel for KDE
-        """
+        """Gaussian kernel for KDE"""
         return (1.0 / np.sqrt(2 * np.pi) / bw) * np.exp(
             -((X - x) ** 2) / (bw ** 2 * 2.0)
         )
@@ -124,7 +115,7 @@ def pz_fn(self, z):
 
 @register_pytree_node_class
 class systematic_shift(redshift_distribution):
-    """ Implements a systematic shift in a redshift distribution
+    """Implements a systematic shift in a redshift distribution
     TODO: Find a better name for this
 
     Arguments:

diff --git a/jax_cosmo/sparse.py b/jax_cosmo/sparse.py
@@ -27,10 +27,6 @@
  - :fun:`sparse_dot_sparse`
  - :fun:`dense_dot_sparse_dot_dense`
 """
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import functools
 
 import jax.numpy as np
@@ -39,14 +35,12 @@
 
 
 def is_sparse(sparse):
-    """Test if the input is interpretable as a sparse matrix.
-    """
+    """Test if the input is interpretable as a sparse matrix."""
     return np.asarray(sparse).ndim == 3
 
 
 def check_sparse(sparse, square=False):
-    """Check for a valid sparse matrix.
-    """
+    """Check for a valid sparse matrix."""
     sparse = np.asarray(sparse)
     if sparse.ndim != 3:
         raise ValueError("Expected 3D array of sparse diagonals.")

diff --git a/jax_cosmo/transfer.py b/jax_cosmo/transfer.py
@@ -1,8 +1,4 @@
 # This module contains various transfer functions from the literatu
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import jax.numpy as np
 
 import jax_cosmo.background as bkgrd
@@ -12,7 +8,7 @@
 
 
 def Eisenstein_Hu(cosmo, k, type="eisenhu_osc"):
-    """ Computes the Eisenstein & Hu matter transfer function.
+    """Computes the Eisenstein & Hu matter transfer function.
 
     Parameters
     ----------

diff --git a/jax_cosmo/utils.py b/jax_cosmo/utils.py
@@ -1,9 +1,4 @@
 # This defines a few utility functions
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-
 def z2a(z):
     """ converts from redshift to scale factor """
     return 1.0 / (1.0 + z)