PyCosmo 2.1.0__tar.gz → 2.1.1__tar.gz

{PyCosmo-2.1.0/src/PyCosmo.egg-info → PyCosmo-2.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: PyCosmo
-Version: 2.1.0
+Version: 2.1.1
 Summary: A multi-purpose cosmology calculation tool
 Author: Uwe Schmitt et al
 Author-email: uwe.schmitt@id.ethz.ch

{PyCosmo-2.1.0 → PyCosmo-2.1.1}/pyproject.toml

@@ -28,7 +28,8 @@ filterwarnings = ["ignore:invalid"]
 legacy_tox_ini = '''
 [tox]
 isolated_build = true
-envlist = py38,py39,py10
+envlist = py39,py310,py311
+package = sdist
 setenv =
     PIP_EXTRA_INDEX_URL=https://cosmo-pypi.phys.ethz.ch/simple
     HTTP_PROXY=http://proxy.ethz.ch:3128
@@ -37,8 +38,8 @@ setenv =
 [testenv]
 skip_install = false
 passenv = *
+allowlist_externals = pytest
 extras = dev
-allowlist_exernals = pytest
 commands =
-    pytest -x -v -ra --color=yes --runslow
+    pytest --runslow -v -ra tests
 '''

{PyCosmo-2.1.0 → PyCosmo-2.1.1}/setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = PyCosmo
-version = 2.1.0
+version = 2.1.1
 license = GPLv3
 description = A multi-purpose cosmology calculation tool
 long_description = file: README.rst
@@ -26,13 +26,13 @@ project_urls =
 [options]
 install_requires = 
 	numpy
-	scipy>=0.14
 	recfast4py>=0.2.3
 	dill
 	matplotlib
-	numba
-	sympy2c>=2
+	numba>=0.57.0rc1
+	sympy2c>=2.2.2
 	numexpr
+	scipy
 packages = find:
 package_dir = 
 	=src
@@ -49,19 +49,25 @@ where = src
 	*.h
 	config/*.ini
 
+[options.entry_points]
+console_scripts = 
+	recompile = PyCosmo:Cosmo.recompile_from_cli
+
 [options.extras_require]
 dev = 
+	Cython
 	twine
-	pytest
+	pytest<7.2.0
 	pytest-cov
 	pytest-regtest
 	tox
 	flake8
-	black @ git+https://github.com/psf/black   # binary wheels are broken on some machines
+	black
 	isort
 	sphinx
 	sphinx-pynpoint-theme
 	myst-parser
+	memray
 docs = 
 	sphinx
 	sphinx-pynpoint-theme

{PyCosmo-2.1.0 → PyCosmo-2.1.1}/setup.py

@@ -1,9 +1,7 @@
 #!/usr/bin/env python
-
-from setuptools import Extension, setup
-from Cython.Build import cythonize
-
 import numpy
+from Cython.Build import cythonize
+from setuptools import Extension, setup
 
 files = [
     "const.c",

{PyCosmo-2.1.0 → PyCosmo-2.1.1}/src/PyCosmo/Background.py

@@ -27,8 +27,7 @@ def _call(wrapper, function, a=None):
     elif isinstance(a, list):
         a = np.array(a)
     assert isinstance(a, np.ndarray)
-    a = a.astype(float)
-    return getattr(wrapper, function)(a)
+    return getattr(wrapper, function)(a.astype(float))
 
 
 class Background(object):
@@ -80,12 +79,11 @@ class Background(object):
     _hubble = H
 
     def _H2_H02_a(self, a=1.0):
-        return (
-            self._H2_H02_Omegar_a(a=a)
-            + self._H2_H02_Omegam_a(a=a)
-            + self._H2_H02_Omegak_a(a=a)
-            + self._H2_H02_Omegal_a(a=a)
-        )
+        result = self._H2_H02_Omegar_a(a=a)
+        result += self._H2_H02_Omegam_a(a=a)
+        result += self._H2_H02_Omegak_a(a=a)
+        result += self._H2_H02_Omegal_a(a=a)
+        return result
 
     def _H2_H02_Omegar_a(self, a=1.0):
         return self._params.omega_r / a**4
@@ -294,7 +292,12 @@ class Background(object):
         """
         a = np.atleast_1d(a)
         return np.array(
-            [integrate.quadrature(self._tau_intgd, np.log(aa), 0.0)[0] for aa in a]
+            [
+                integrate.quadrature(
+                    self._tau_intgd, np.log(aa), 0.0, rtol=1e-6, maxiter=100
+                )[0]
+                for aa in a
+            ]
         )
 
     def _tau_intgd(self, lna):
@@ -435,7 +438,9 @@ class Background(object):
         input: a - scale factor [1]
         output: [1]
         """
-        return self._H2_H02_Omegal_a(a=a) / self._H2_H02_a(a=a)
+        result = self._H2_H02_Omegal_a(a=a)
+        result /= self._H2_H02_a(a=a)
+        return result
 
     def _omega_nu_m_a(self, a=1.0):
         """

{PyCosmo-2.1.0 → PyCosmo-2.1.1}/src/PyCosmo/BoltzmannSolver/LinearPerturbationBoltzmann.py

@@ -16,7 +16,9 @@
 
 import itertools
 import warnings
+from datetime import datetime
 from functools import lru_cache
+from multiprocessing import current_process
 
 import numpy as np
 
@@ -24,9 +26,8 @@ from ..LinearPerturbationBase import LinearPerturbationBase
 from .Fields import Fields
 from .LsodaSolver import LsodaSolver
 
-
 # use decorator to print messages only once per model:
-@lru_cache(maxsize=1)
+@lru_cache(maxsize=1)  # noqa: E302
 def print_new_traces(cache_file):
     print()
     print("new traces detected! you might want to run")
@@ -34,6 +35,39 @@ def print_new_traces(cache_file):
     print()
 
 
+proc_name = current_process().name
+
+
+def trace_call(frame, event, arg):
+    if proc_name.endswith("-2"):
+        return trace_call
+    if event in ("line",):
+        return trace_call
+    fname = frame.f_code.co_filename
+    if any(
+        n in fname
+        for n in (
+            "pprint",
+            "dill",
+            "pickle",
+            "numba",
+            "logging",
+            "ini_handling",
+            "llvmlite",
+            "traceback",
+            "weakref",
+        )
+    ):
+        return trace_call
+    print(datetime.now(), end=" ")
+    fr = frame
+    print(proc_name, f"{event:8s}", end="")
+    while (fr := fr.f_back) is not None:
+        print("  ", end="")
+    print(fname, frame.f_code.co_name, frame.f_lineno, flush=True)
+    return trace_call
+
+
 class LinearPerturbationBoltzmann(LinearPerturbationBase):
     """
     Class for computing linear perturbations by solving the Einstein-Boltzmann ODE
@@ -57,6 +91,7 @@ class LinearPerturbationBoltzmann(LinearPerturbationBase):
             self._params.boltzmann_atol,
             self._params.boltzmann_max_bdf_order,
             self._params.boltzmann_max_iter,
+            self._params.boltzmann_h0,
             self._cosmo,
             self._params.fast_solver,
         )
@@ -152,17 +187,17 @@ class LinearPerturbationBoltzmann(LinearPerturbationBase):
             # using equation in section 2.4 of notes
             norm = (
                 2.0
-                * np.pi ** 2
-                * self._params.pk_norm ** 2
+                * np.pi**2
+                * self._params.pk_norm**2
                 * (self._params.c / self._params.H0) ** (3.0 + self._params.n)
             )
 
             if diag_only:
-                pk = norm * growth ** 2 * k ** self._params.n * T_k ** 2
+                pk = norm * growth**2 * k**self._params.n * T_k**2
             else:
                 # pk = norm * np.outer(growth.T ** 2, k ** self._params.n * T_k ** 2).T
                 pk = (
-                    norm * growth ** 2 * (k ** self._params.n * T_k ** 2).reshape(-1, 1)
+                    norm * growth**2 * (k**self._params.n * T_k**2).reshape(-1, 1)
                 )
 
         if (self._params.pk_norm_type == "A_s") or (
@@ -178,7 +213,7 @@ class LinearPerturbationBoltzmann(LinearPerturbationBase):
     def _a_s_powerspec(self, a, k, diag_only):
         norm = (
             2.0
-            * np.pi ** 2
+            * np.pi**2
             * self._params.As_norm
             / self._params.k_pivot ** (self._params.n - 1)
         )
@@ -191,26 +226,76 @@ class LinearPerturbationBoltzmann(LinearPerturbationBase):
         if self._cosmo._pool is None:
             if diag_only:
                 for ai, ki in zip(a, kk):
-                    src_m.append(src_m_column_j((self, ai, ki)))
+                    src_m.append(
+                        src_m_column_j(
+                            (
+                                self._compute_fields_,
+                                self._solver.solve,
+                                self._cosmo._cache_file,
+                                self._params,
+                                self._background.H(ai),
+                                self._wrapper,
+                                ai,
+                                ki,
+                            )
+                        )
+                    )
             else:
                 for ki in kk:
-                    src_m.append(src_m_column_j((self, a, ki)))
+                    src_m.append(
+                        src_m_column_j(
+                            (
+                                self._compute_fields_,
+                                self._solver.solve,
+                                self._cosmo._cache_file,
+                                self._params,
+                                self._background.H(a),
+                                self._wrapper,
+                                a,
+                                ki,
+                            )
+                        )
+                    )
         else:
             kk = np.array(kk)
             perm = np.argsort(kk)
             kk = kk[perm]
             src_m = [None] * len(kk)
             if diag_only:
-                args = [(self, ai, ki) for ai, ki in zip(a, kk)]
+                args = [
+                    (
+                        self._compute_fields_,
+                        self._solver.solve,
+                        self._cosmo._cache_file,
+                        self._params,
+                        self._background.H(ai),
+                        self._wrapper,
+                        ai,
+                        ki,
+                    )
+                    for ai, ki in zip(a, kk)
+                ]
             else:
-                args = [(self, a, ki) for ki in kk]
+                args = [
+                    (
+                        self._compute_fields_,
+                        self._solver.solve,
+                        self._cosmo._cache_file,
+                        self._params,
+                        self._background.H(a),
+                        self._wrapper,
+                        a,
+                        ki,
+                    )
+                    for ki in kk
+                ]
 
             src_m_cols = list(self._cosmo._pool.map(src_m_column_j, args))
             for i, pi in enumerate(perm):
                 src_m[pi] = src_m_cols[i]
 
         src_m = np.vstack(src_m).T
-        return norm * src_m.T ** 2 * k[:, None] ** (self._params.n - 4.0)
+        return norm * src_m.T**2 * k[:, None] ** (self._params.n - 4.0)
 
     def powerspec_cb_a_k(self, a=1.0, k=0.1, diag_only=False):
         """
@@ -239,7 +324,7 @@ class LinearPerturbationBoltzmann(LinearPerturbationBase):
 
         norm = (
             2.0
-            * np.pi ** 2
+            * np.pi**2
             * self._params.pk_norm
             / self._params.k_pivot ** (self._params.n - 1)
         )
@@ -257,10 +342,10 @@ class LinearPerturbationBoltzmann(LinearPerturbationBase):
             theta_m_tot = self._params.omega_dm * theta + self._params.omega_b * theta_b
 
             tk[:, j] = (
-                delta_m_tot + (3 * a * H / onek ** 2) * theta_m_tot
+                delta_m_tot + (3 * a * H / onek**2) * theta_m_tot
             ) / self._params.omega_m
 
-        pk = norm * tk.T ** 2 * k[:, None] ** (self._params.n - 4.0)
+        pk = norm * tk.T**2 * k[:, None] ** (self._params.n - 4.0)
 
         if diag_only:
             return np.diag(pk)
@@ -419,9 +504,32 @@ class LinearPerturbationBoltzmann(LinearPerturbationBase):
         enable_fast_solver=True,
         enable_sparse_lu_solver=True,
     ):
+        return LinearPerturbationBoltzmann._compute_fields_(
+            self._solver.solve,
+            k,
+            grid,
+            sec_factor,
+            keep_lna0,
+            initial_conditions,
+            enable_fast_solver,
+            enable_sparse_lu_solver,
+            self._cosmo._cache_file,
+        )
 
+    @staticmethod
+    def _compute_fields_(
+        solve,
+        k,
+        grid,
+        sec_factor=3,
+        keep_lna0=False,
+        initial_conditions=None,
+        enable_fast_solver=True,
+        enable_sparse_lu_solver=True,
+        cache_file=None,
+    ):
         grid = np.atleast_1d(grid)
-        grid, solver_result, meta = self._solver.solve(
+        grid, solver_result, meta = solve(
             k,
             grid,
             sec_factor,
@@ -434,7 +542,7 @@ class LinearPerturbationBoltzmann(LinearPerturbationBase):
         if (isinstance(new_traces, dict) and new_traces) or (
             isinstance(new_traces, list) and any(new_traces)
         ):
-            print_new_traces(self._cosmo._cache_file)
+            print_new_traces(cache_file)
         return grid, solver_result, meta, keep_lna0
 
     def P_R_k(self, k):
@@ -458,10 +566,12 @@ class LinearPerturbationBoltzmann(LinearPerturbationBase):
         return norm * k ** (self._params.n - 1)
 
 
-def src_m_column_j(arg):
-    self, a, onek = arg
+def src_m_column_j(args):
+    compute_fields_, solve, cache_file, _params, H_, _wrapper, a, onek = args
     a = np.atleast_1d(a)
-    grid, y, meta, keep_lna0 = self._compute_fields(onek, grid=np.log(a))
+    grid, y, meta, keep_lna0 = compute_fields_(
+        solve, onek, grid=np.log(a), cache_file=cache_file
+    )
 
     n_grid = len(grid)
     a = a[-n_grid:]
@@ -469,34 +579,34 @@ def src_m_column_j(arg):
     theta = y[:, 2] * onek
     delta_b = y[:, 3]
     theta_b = y[:, 4] * onek
-    H = self._background.H(a) / (self._params.H0 * self._params.rh)
+    H = H_ / (_params.H0 * _params.rh)
     omega_nu_m = 0.0
     P_nu_m = 0.0
-    omega_m_tot = self._params.omega_m
-    omega_plus_p_m_tot = self._params.omega_m
+    omega_m_tot = _params.omega_m
+    omega_plus_p_m_tot = _params.omega_m
     # include massive neutrinos in the total matter term
     src_col = np.zeros((n_grid,))
-    if self._params.N_massive_nu != 0:
-        omega_nu_m = self._wrapper.omega_nu_m_ufunc(a)
-        P_nu_m = self._wrapper.P_nu_m_ufunc(a)
+    if _params.N_massive_nu != 0:
+        omega_nu_m = _wrapper.omega_nu_m_ufunc(a)
+        P_nu_m = _wrapper.P_nu_m_ufunc(a)
         omega_m_tot += omega_nu_m / a
         omega_plus_p_m_tot += (omega_nu_m + P_nu_m) / a
-        delta_nu_m = self._wrapper.delta_nu_m_ufunc(a, y[1:])
-        theta_nu_m = onek * self._wrapper.u_nu_m_ufunc(a, y[1:])
+        delta_nu_m = _wrapper.delta_nu_m_ufunc(a, y[1:])
+        theta_nu_m = onek * _wrapper.u_nu_m_ufunc(a, y[1:])
         src_col += (
             delta_nu_m * omega_nu_m / a / omega_m_tot
-            + (3 * a * H / onek ** 2)
+            + (3 * a * H / onek**2)
             * theta_nu_m
             * (omega_nu_m + P_nu_m)
             / a
             / omega_plus_p_m_tot
         )
 
-    delta_m_tot = self._params.omega_dm * delta + self._params.omega_b * delta_b
-    theta_m_tot = self._params.omega_dm * theta + self._params.omega_b * theta_b
+    delta_m_tot = _params.omega_dm * delta + _params.omega_b * delta_b
+    theta_m_tot = _params.omega_dm * theta + _params.omega_b * theta_b
 
     src_col += (
         delta_m_tot / omega_m_tot
-        + (3 * a * H / onek ** 2) * theta_m_tot / omega_plus_p_m_tot
+        + (3 * a * H / onek**2) * theta_m_tot / omega_plus_p_m_tot
     )
     return src_col

{PyCosmo-2.1.0 → PyCosmo-2.1.1}/src/PyCosmo/BoltzmannSolver/LsodaSolver.py

@@ -33,6 +33,7 @@ class LsodaSolver:
         atol,
         max_bdf_order,
         max_iter,
+        h0,
         cosmology,
         fast_solver,
         **extra_args_for_solver,
@@ -42,15 +43,16 @@ class LsodaSolver:
         self.atol = atol
         self.max_bdf_order = max_bdf_order
         self.max_iter = max_iter
+        self.h0 = h0
         self.extra_args_for_solver = extra_args_for_solver
 
         self.params = cosmology.params
-        self.cosmology = cosmology
-        self.background = cosmology.background
-        self.fast_solver = fast_solver
-
+        # self.background = cosmology.background
         self.wrapper = cosmology._wrapper
 
+        assert fast_solver, "slow solver not supported anymore"
+        self.fast_solver = fast_solver
+
     def initial_conditions(self, k):
         self.wrapper.set_globals(k=k)
 
@@ -59,8 +61,8 @@ class LsodaSolver:
         )
         if initial_conditions is None:
             raise ValueError(
-                "initial conditions '{}' not implemented.".format(
-                    self.params.initial_conditions
+                "{}: initial conditions '{}' not implemented.".format(
+                    self.wrapper, self.params.initial_conditions
                 )
             )
 
@@ -108,27 +110,20 @@ class LsodaSolver:
         grid = grid[grid > lna0]
         grid = np.concatenate(([lna0], grid))
 
-        atol = self.atol
-        rtol = self.rtol
-        max_order = self.max_bdf_order
-        max_iter = self.max_iter
-
-        if self.fast_solver:
-            solver = getattr(self.wrapper, "solve_fast_" + self.SOLVER_ID)
-        else:
-            solver = getattr(self.wrapper, "solve_" + self.SOLVER_ID)
+        solver = getattr(self.wrapper, "solve_fast_" + self.SOLVER_ID)
 
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")
             y, meta = solver(
                 y_0,
                 grid,
-                rtol,
-                atol,
-                max_order=max_order,
-                max_iter=max_iter,
+                self.rtol,
+                self.atol,
+                max_order=self.max_bdf_order,
+                max_iter=self.max_iter,
                 enable_fast_solver=enable_fast_solver,
                 enable_sparse_lu_solver=enable_sparse_lu_solver,
+                h0=self.h0,
             )
         if meta["istate"] <= 0:
             raise RuntimeError(f"solving ode failed for k = {k}, meta={meta}")