Source code for camb.symbolic

"""

This module defines the scalar linear perturbation equations for standard LCDM cosmology, using sympy.
It uses the covariant perturbation notation, but includes functions to project into the
Newtonian or synchronous gauge, as well as constructing general gauge invariant quantities.
It uses "t" as the conformal time variable (=tau in the fortran code).

For a guide to usage and content see the `ScalEqs notebook <https://camb.readthedocs.io/en/latest/ScalEqs.html>`_

As well as defining standard quantities, and how they map to CAMB variables, there are also functions for
converting a symbolic expression to CAMB source code, and compiling custom sources for use with CAMB
(as used by :meth:`.model.CAMBparams.set_custom_scalar_sources`, :meth:`.results.CAMBdata.get_time_evolution`)

A Lewis July 2017
"""

import ctypes
import os
import sympy
from sympy import diff, Eq, simplify, Function, Symbol
from sympy.core.relational import Relational
from sympy.abc import t, kappa, K, k

# Background variables

tau0 = Symbol('tau0', description='conformal time today')
tau_maxvis = Symbol('tau_maxvis', description='conformal time of peak visibility')

f_K = Function('f_K', description='comoving angular diameter distance')

H = Function('H', description='comoving hubble parameter', camb_var='adotoa')(t)
rho = Function('rho', description='total density', camb_var='grho', camb_sub='grho/kappa/a**2')(t)
P = Function('P', description='total pressure', camb_var='gpres', camb_sub='gpres/kappa/a**2')(t)
a = Function('a', description='scale factor')(t)

rho_b = Function('rho_b', description='baryon density', camb_var='grhob_t', camb_sub='grhob_t/kappa/a**2')(t)
rho_c = Function('rho_c', description='CDM density', camb_var='grhoc_t', camb_sub='grhoc_t/kappa/a**2')(t)
rho_g = Function('rho_g', description='photon density', camb_var='grhog_t', camb_sub='grhog_t/kappa/a**2')(t)
rho_r = Function('rho_r', description='massless neutrino density', camb_var='grhor_t', camb_sub='grhor_t/kappa/a**2')(t)
rho_nu = Function('rho_nu', description='massive neutrino density', camb_var='grhonu_t',
                  camb_sub='grhonu_t/kappa/a**2')(t)
rho_de = Function('rho_de', description='dark energy density', camb_var='grhov_t', camb_sub='grhov_t/kappa/a**2')(t)

p_b = Function('p_b', description='baryon pressure', camb_sub='0')(t)
p_nu = Function('p_nu', description='massive neutrino pressure')(t)
w_de = Function('w_de', camb_var='w_lam', description='fluid dark energy equation of state')(t)
p_g = rho_g / 3
p_r = rho_r / 3
p_c = 0
p_de = w_de * rho_de

opacity = Function('opacity', description='opacity, a n_e sigma_t')(t)
visibility = Function('visibility', description='ionization visibility')(t)
exptau = Function('exptau', description='exp(-tau)')(t)


def subs(eqs, expr):
    # generalization to act on lists of equations, and support using Eq equations.
    # lists and dicts are substituted irrespective of order, so no RHS variables
    # are substituted by other elements
    if isinstance(expr, (list, tuple)):
        res = [subs(eqs, ex) for ex in expr]
        return [x for x in res if x is not True]
    if not isinstance(expr, (sympy.Expr, Relational)):
        return expr
    if isinstance(eqs, dict):
        return expr.subs(eqs, simultaneous=True)
    else:
        if not isinstance(eqs, (list, tuple)):
            return expr.subs(eqs.lhs, eqs.rhs)
        eqs = [(eq.lhs, eq.rhs) for eq in eqs]
        return expr.subs(eqs, simultaneous=True)


def solve(eq, x):
    res = sympy.solve(eq, x)
    if len(res) == 1:
        return res[0]
    else:
        return res


half = sympy.Rational(1, 2)
third = sympy.Rational(1, 3)
Kf = sympy.IndexedBase('Kf', shape=(sympy.oo,))
K_fac = Symbol('Kf_1')  # currently sympy bugs just using Kf[1], so treat as separate symbol
K_fac_sub = 1 - 3 * K / k ** 2
K_sub = Eq(K, solve(K_fac_sub - K_fac, K))

H_t = diff(a, t) / a
dH = -a ** 2 / 6 * kappa * (rho + 3 * P)
Friedmann = Eq(H ** 2, a ** 2 * kappa * rho / 3 - K)
Friedmann_subs = [Friedmann, Eq(diff(H, t), dH), Eq(diff(a, t), a * H)]
Friedmann_Kfac_subs = subs(K_sub, Friedmann_subs)

delta_frame = Function('uprime')(t)


[docs] def LinearPerturbation(name, species=None, camb_var=None, camb_sub=None, frame_dependence=None, description=None): """ Returns as linear perturbation variable, a function of conformal time t. Use help(x) to quickly view all quantities defined for the result. :param name: sympy name for the Function :param species: tag for the species if relevant (not used) :param camb_var: relevant CAMB fortran variable :param camb_sub: if not equal to camb_var, and string giving the expression in CAMB variables :param frame_dependence: the change in the perturbation when the frame 4-velocity u change from u to u + delta_frame. Should be a numpy expression involving delta_frame. :param description: string describing variable :return: sympy Function instance (function of t), with attributes set to the arguments above. """ if isinstance(camb_var, list): camb_var = tuple(camb_var) f = Function(name, species=species, camb_var=camb_var, camb_sub=camb_sub, perturbation_order=1, frame_dependence=frame_dependence, description=description) return f(t)
def list_perturbations(expr, lst=None): if lst is None: lst = [] if getattr(expr, 'perturbation_order', None) and expr not in lst: lst.append(expr) for arg in expr.args: list_perturbations(arg, lst) return lst def list_frame_dependent_vars(expr, lst=None): if lst is None: lst = [] if getattr(expr, 'frame_dependence', None) and expr not in lst: lst.append(expr) for arg in expr.args: list_frame_dependent_vars(arg, lst) return lst def simplify_sum(expr): if isinstance(expr, sympy.Add): return sympy.Add(*[simplify(term) for term in expr.args]) else: return expr def frame_change(expr, delta_u=None, total=False): if getattr(expr, 'frame_dependence', None): res = expr.frame_dependence if delta_u is not None: res = res.subs(delta_frame, delta_u) if total: res += expr return res else: if isinstance(expr, (list, tuple)): return [frame_change(x) for x in expr] perts = list_frame_dependent_vars(expr) res = subs([Eq(pert, pert + pert.frame_dependence) for pert in perts], expr) if delta_u is not None and isinstance(res, sympy.Expr): res = res.subs(delta_frame, delta_u) if not total: res -= expr res = simplify(res).expand() return simplify_sum(res.collect(list_frame_dependent_vars(res))) # Perturbation variables # gauge-invariant potential phi = LinearPerturbation('phi', description='Weyl potential') eta = LinearPerturbation('eta', camb_var='etak', camb_sub='-2*etak/k', description='three-curvature', frame_dependence=-2 * K_fac * H * delta_frame / k) sigma = LinearPerturbation('sigma', description='shear', frame_dependence=delta_frame) A = LinearPerturbation('A', camb_sub=0, description='acceleration', frame_dependence=(diff(delta_frame, t) + H * delta_frame) / k) z = LinearPerturbation('z', description='expansion rate perturbation', # frame_dependence = delta_frame - 3*(dH- H**2)*delta_frame/k**2) frame_dependence=K_fac * delta_frame + 3 * kappa * a ** 2 * (rho + P) * delta_frame / ( 2 * k ** 2)) hdot = LinearPerturbation('hdot', camb_sub=k / 3 * z, description='time derivative of scale factor perturbation', frame_dependence=k * delta_frame / 3 - diff(delta_frame * H, t) / k) delta = LinearPerturbation('delta', camb_var='dgrho', camb_sub='dgrho/kappa/a**2', description='total density perturbation', frame_dependence=3 * H * (rho + P) * delta_frame / k) delta_P = LinearPerturbation('delta_P', camb_sub='error', description='total pressure perturbation', frame_dependence=-diff(P, t) * delta_frame / k) q = LinearPerturbation('q', camb_var='dgq', camb_sub='dgq/kappa/a**2', description='total heat flux', frame_dependence=-(rho + P) * delta_frame) Pi = LinearPerturbation('Pi', camb_var='dgpi', camb_sub='dgpi/a**2/kappa', description='total anisotropic stress') # quadrupole source polter = Function('polter')(t) # Newtonian gauge variables (in general equal to gauge invariant potentials) Phi_N = LinearPerturbation('Phi_N', description='Newtonian gauge curvature potential') Psi_N = LinearPerturbation('Psi_N', description='Newtonian gauge acceleration potential') # Synchronous gauge variables eta_s = LinearPerturbation('eta_s', description='Synchronous gauge curvature variable') hdot_s = LinearPerturbation('hdot_s', description='Synchronous gauge diff(h,t)') # velocities, and dimensionless heat flux [(rho_i+P_i)v_i = rho_i q_i] q_r = LinearPerturbation('q_r', species='r', camb_var='qr', description='massless neutrino heat flux', frame_dependence=-4 * delta_frame / 3) q_g = LinearPerturbation('q_g', species='g', camb_var='qg', description='photon heat flux', frame_dependence=-4 * delta_frame / 3) q_nu = LinearPerturbation('q_nu', species='nu', camb_var='qnu', description='massive neutrino heat flux', frame_dependence=-(1 + p_nu / rho_nu) * delta_frame) v_c = LinearPerturbation('v_c', species='c', description='CDM velocity', frame_dependence=-delta_frame) v_b = LinearPerturbation('v_b', species='b', camb_var='vb', description='baryon velocity', frame_dependence=-delta_frame) v_de = LinearPerturbation('v_de', species='de', camb_var=['qde', 'w_lam'], camb_sub='qde/(1+w_lam)', description='dark energy velocity', frame_dependence=-delta_frame) Delta_b = LinearPerturbation('Delta_b', species='b', camb_var='clxb', description='fractional baryon density perturbation', frame_dependence=3 * H * (1 + p_b / rho_b) * delta_frame / k) Delta_P_b = LinearPerturbation('Delta_P_b', camb_sub='delta_p_b', description='fractional baryon pressure perturbation', frame_dependence=-diff(p_b, t) / rho_b * delta_frame / k) Delta_c = LinearPerturbation('Delta_c', species='c', camb_var='clxc', description='fractional CDM density perturbation', frame_dependence=3 * H * delta_frame / k) Delta_r = LinearPerturbation('Delta_r', species='r', camb_var='clxr', description='fractional massless neutrino density perturbation', frame_dependence=4 * H * delta_frame / k) Delta_nu = LinearPerturbation('Delta_nu', species='nu', camb_var='clxnu', description='fractional massive neutrino density perturbation', frame_dependence=3 * H * (1 + p_nu / rho_nu) * delta_frame / k) Delta_P_nu = LinearPerturbation('Delta_P_nu', camb_sub='error', description='fractional massive neutrino pressure perturbation', frame_dependence=-diff(p_nu, t) / rho_nu * delta_frame / k) Delta_g = LinearPerturbation('Delta_g', species='g', camb_var='clxg', description='fractional CDM density perturbation', frame_dependence=4 * H * delta_frame / k) Delta_de = LinearPerturbation('Delta_de', species='de', camb_var='clxde', description='fractional dark energy density perturbation', frame_dependence=3 * H * (1 + w_de) * delta_frame / k) # Sound speeds csq_b = LinearPerturbation('c_sb^2', species='b', camb_var=['delta_p_b', 'clxb'], camb_sub='delta_p_b/clxb', description='baryon sound speed') csq_b.frame_dependence = (csq_b * Delta_b - diff(p_b, t) / rho_b * delta_frame / k) / ( Delta_b + 3 * H * (1 + p_b / rho_b) * delta_frame / k) - csq_b # dark energy sound speed defined in rest frame so gauge invariant csqhat_de = LinearPerturbation('chat_sde^2', species='de', camb_var='cs2_lam', description='rest frame dark energy sound speed') # Anisotropic stress pi_g = LinearPerturbation('pi_g', species='g', camb_var='pig', description='photon anisotropic stress') pi_r = LinearPerturbation('pi_r', species='r', camb_var='pir', description='massless neutrino anisotropic stress') pi_nu = LinearPerturbation('pi_nu', species='r', camb_var='pinu', description='massive neutrino anisotropic stress') E_2 = LinearPerturbation('E_2', species='E', description='E polarization quadrupole') E_3 = LinearPerturbation('E_3', species='E', description='E polarization octopole') J_3 = LinearPerturbation('J_3', species='g', description='photon octopole') polter_t = sympy.Rational(2, 15) * (3 * pi_g / 4 + 9 * E_2 / 2) polter_sub = Eq(polter, polter_t) background_eqs = [ Eq(diff(a, t), H * a), Eq(diff(H, t), dH), Eq(diff(exptau, t), visibility) ] # constraint equations relating variables cons1 = k ** 3 * K_fac * phi + kappa / 2 * a ** 2 * k * (delta + Pi * K_fac) + 3 * kappa / 2 * a ** 2 * H * q cons2 = k ** 2 * eta - kappa * a ** 2 * delta + 2 * k * H * z cons3 = 2 * third * (k / a) ** 2 * (z - K_fac * sigma) + kappa * q cons4 = hdot - k / 3 * z + H * A constraints = [cons1, cons2, cons3, cons4] def constraint_subs_for_variable_set(variables=(z, sigma, phi, hdot)): return solve(constraints, variables) # Various substitutions for variables in terms of other variables (using constraints) var_subs = constraint_subs_for_variable_set() hdot_sub, phi_sub, sigma_sub, z_sub = [Eq(variable, subs(var_subs, variable)) for variable in [hdot, phi, sigma, z]] q_sub = Eq(q, subs(Eq(z, subs(var_subs, z)), solve(cons3, q))) # Evolution equations dz = -H * z - kappa * a ** 2 / k / 2 * (delta + 3 * delta_P) + (3 * kappa * a ** 2 * ( rho + P) / 2) * A / k + k * K_fac * A dsigma = -H * sigma + k * (phi + A) - half * kappa * a ** 2 / k * Pi deta = -1 / k * (2 * K * z + kappa * a ** 2 * q + 2 * K_fac * k * H * A) dphi = -H * phi + half / k ** 2 * kappa * a ** 2 * (k * (rho + P) * sigma + k * q - diff(Pi, t) - H * Pi) # Alternative equation deta_2 = Eq(diff(eta, t), K_fac * (2 * hdot - 2 * k / 3 * sigma)) pert_eqs = [ Eq(diff(z, t), dz), Eq(diff(sigma, t), dsigma), Eq(diff(eta, t), deta), Eq(diff(phi, t), dphi)] drag_t = opacity * (4 * v_b / 3 - q_g) total_eqs = [ Eq(diff(q, t), k * delta_P - 4 * H * q - 2 * third * k * K_fac * Pi - (rho + P) * k * A), Eq(diff(delta, t), -k * q - 3 * H * (delta + delta_P) - 3 * hdot * (rho + P)) ] tot_eqs = total_eqs + pert_eqs + background_eqs # frame names specify a frame where a certain variable is zero frame_names = {'CDM': v_c, 'Newtonian': sigma, 'comoving': q, 'flat': eta, 'constant density': delta}
[docs] def make_frame_invariant(expr, frame='CDM'): """ Makes the quantity gauge invariant, assuming currently evaluated in frame 'frame'. frame can either be a string frame name, or a variable that is zero in the current frame, e.g. frame = Delta_g gives the constant photon density frame. So make_frame_invariant(sigma, frame=Delta_g) will return the combination of sigma and Delta_g that is frame invariant (and equal to just sigma when Delta_g=0). """ if isinstance(frame, str): if frame not in frame_names: raise ValueError('Unknown frame names: %s' % frame) frame = frame_names[frame] if isinstance(expr, Eq): return simplify( Eq(make_frame_invariant(expr.lhs, frame), make_frame_invariant(expr.rhs, frame))) if isinstance(expr, (list, tuple)): return [make_frame_invariant(x, frame) for x in expr] # do frame change to make frame variable zero # special case of frame=A, equivalent to v_c frame by evolution equation if frame == A: frame = v_c delta_u = solve(frame_change(frame, total=True), delta_frame) if delta_frame in delta_u.atoms(Function): raise ValueError( 'Cannot solve for change of frame. Currently only supports algebraic frame changes: %s' % delta_u) return frame_change(expr, delta_u=delta_u, total=True)
# Things for relating to Newtonian (zero shear) gauge # Newtonian gauge variables in general. Note sigma here is the shear, Pi the anisotropic stress Newt_vars = [Eq(Psi_N, -A + (diff(sigma, t) + H * sigma) / k), Eq(Phi_N, -eta / 2 / K_fac - H * sigma / k)] Newtonian_var_subs = [Eq(Phi_N, phi + half * a ** 2 * kappa * Pi / k ** 2), Eq(Psi_N, phi - half * a ** 2 * kappa * Pi / k ** 2)] Newtonian_subs = [Eq(A, -Psi_N), Eq(diff(sigma, t, t), 0), Eq(diff(sigma, t), 0), Eq(sigma, 0), Eq(phi, (Phi_N + Psi_N) / 2), Eq(eta, -2 * Phi_N * K_fac), Eq(hdot, -diff(Phi_N, t))] Newtonian_subs += [Eq(z, subs(Newtonian_subs, solve(cons4, z)))]
[docs] def newtonian_gauge(x): r""" Evaluates an expression in the Newtonian gauge (zero shear, sigma=0). Converts to using conventional metric perturbation variables for metric .. math:: ds^2 = a^2\left( (1+2\Psi_N)d t^2 - (1-2\Phi_N)\delta_{ij}dx^idx^j\right) :param x: expression :return: expression evaluated in the Newtonian gauge """ if isinstance(x, (list, tuple)): res = [newtonian_gauge(y) for y in x] return [x for x in res if x is not True] res = subs(Newtonian_subs, x) if isinstance(res, sympy.Expr): res = simplify(res.doit()) res2 = simplify(res.subs(Psi_N, Phi_N - a ** 2 * kappa * Pi / k ** 2)) if len(str(res2)) < len(str(res)): res = res2 return res
# Things for relating to cdm frame or traditional synchronous gauge variables # In the "cdm frame" we just fix to zero acceleration, comoving with CDM # the covariant hdot and eta variables differ from Ma & Bertshinger by factors (the latter have _s after the name) cdm_subs = [Eq(diff(A, t), 0), Eq(A, 0), Eq(diff(v_c, t), 0), Eq(v_c, 0)]
[docs] def cdm_gauge(x): r""" Evaluates an expression in the CDM frame :math:`(v_c=0, A=0)`. Equivalent to the synchronous gauge but using the covariant variable names. :param x: expression :return: expression evaluated in CDM frame. """ if isinstance(x, (list, tuple)): return [cdm_gauge(y) for y in x] return simplify(subs(cdm_subs, x))
synchronous_vars = [Eq(eta_s, -make_frame_invariant(eta, v_c) / 2 / K_fac), Eq(hdot_s, 6 * (hdot + H * A + v_c / k * (k ** 2 / 3 * K_fac + kappa * a ** 2 / 2 * (rho + P))))] synchronous_subs = [Eq(eta, -2 * K_fac * eta_s), Eq(hdot, hdot_s / 6), Eq(phi, subs(var_subs, phi))] synchronous_subs += [subs(synchronous_subs, Eq(z, solve(cons4, z))), Eq(sigma, (hdot_s + diff(6 * eta_s, t)) / 2 / k)] synchronous_subs = cdm_gauge(synchronous_subs) + cdm_subs
[docs] def synchronous_gauge(x): """ evaluates an expression in the synchronous gauge, using conventional synchronous-gauge variables. :param x: expression :return: synchronous gauge variable expression """ if isinstance(x, (list, tuple)): res = [synchronous_gauge(y) for y in x] return [x for x in res if x is not True] res = subs(synchronous_subs, x) if isinstance(res, sympy.Expr): return simplify(res.doit()) return simplify(res)
# Fluid components density_eqs = [ Eq(diff(rho_b, t), -3 * H * (rho_b + p_b)), Eq(diff(rho_c, t), -3 * H * rho_c), Eq(diff(rho_g, t), -4 * H * rho_g), Eq(diff(rho_r, t), -4 * H * rho_r), Eq(diff(rho_nu, t), -3 * H * (rho_nu + p_nu)), Eq(diff(rho_de, t), -3 * H * (rho_de * (1 + w_de))) ] delta_eqs = [ Eq(diff(Delta_r, t), -4 * hdot - k * q_r), Eq(diff(Delta_g, t), -4 * hdot - k * q_g), Eq(diff(Delta_b, t), -(1 + p_b / rho_b) * (3 * hdot + k * v_b) + (p_b / rho_b - csq_b) * 3 * H * Delta_b), Eq(diff(Delta_c, t), -3 * hdot - k * v_c), Eq(diff(Delta_nu, t), -3 * (1 + p_nu / rho_nu) * hdot - k * q_nu + 3 * H * (-Delta_P_nu + Delta_nu * p_nu / rho_nu)), Eq(diff(Delta_de, t), -3 * (1 + w_de) * hdot - (1 + w_de) * k * v_de - 3 * H * (csqhat_de - w_de) * ( Delta_de + 3 * H * (1 + w_de) * v_de / k) - 3 * H * diff(w_de, t) * v_de / k) ] vel_eqs = [ Eq(diff(q_r, t), -2 * third * k * pi_r * K_fac + third * k * Delta_r - 4 * k / 3 * A), Eq(diff(q_g, t), -2 * third * k * pi_g * K_fac + third * k * Delta_g + drag_t - 4 * k / 3 * A), Eq(diff(v_b, t), -k * A - H * v_b - diff(p_b, t) * v_b / (rho_b + p_b) + 1 / (rho_b + p_b) * (rho_b * csq_b * k * Delta_b - rho_g * drag_t)), Eq(diff(v_c, t), -H * v_c - k * A), Eq(diff(q_nu, t), -H * (1 - 3 * p_nu / rho_nu) * q_nu - k / 3 * (2 * K_fac * pi_nu - 3 * Delta_P_nu) - (1 + p_nu / rho_nu) * k * A), Eq(diff(v_de, t), k * csqhat_de * Delta_de / (1 + w_de) - H * (1 - 3 * csqhat_de) * v_de - k * A) ] component_eqs = density_eqs + delta_eqs + vel_eqs rho_t = rho_b + rho_c + rho_r + rho_g + rho_nu + rho_de P_t = third * (rho_r + rho_g) + p_b + p_nu + w_de * rho_de tot_subs = [ Eq(rho, rho_t), Eq(P, P_t) ] # Note that csqhat_de is defined in the dark energy rest-frame, # so this is general-gauge result for pressure perturbation: Delta_P_de = ( csqhat_de * Delta_de + 3 * H * v_de / k * (1 + w_de) * (csqhat_de - w_de + diff(w_de, t) / 3 / H / (1 + w_de))) tot_pert_subs = [ Eq(Pi, rho_g * pi_g + rho_r * pi_r + rho_nu * pi_nu), Eq(delta, rho_g * Delta_g + rho_r * Delta_r + rho_b * Delta_b + rho_c * Delta_c + rho_nu * Delta_nu + rho_de * Delta_de), Eq(q, rho_g * q_g + rho_r * q_r + rho_c * v_c + (rho_b + p_b) * v_b + rho_nu * q_nu + rho_de * (1 + w_de) * v_de), Eq(delta_P, rho_nu * Delta_P_nu + third * (rho_g * Delta_g + rho_r * Delta_r) + csq_b * Delta_b * rho_b + Delta_P_de * rho_de) ] def define_variable(name, namespace=None, order=1): namespace = namespace or globals() if name not in namespace: namespace[name] = sympy.Function(name, perturbation_order=order)(t) return namespace[name] def define_variables(names, namespace=None, order=1): return [define_variable(name, namespace, order) for name in names.split()] def _make_index_func(name, ell, namespace=None): name += '_' + str(ell) return define_variable(name, namespace) # Boltzmann hierarchies. Haven't included massive neutrinos here as yet. def J_eq(L): # photons assert (L > 1) Gl = _make_index_func('J', L) Glp = _make_index_func('J', L + 1) Glm = _make_index_func('J', L - 1) eq = -k / (2 * L + 1) * ((L + 1) * Kf[L] * Glp - L * Glm) - opacity * Gl if L == 2: eq = eq + 8 * k / 15 * sigma + opacity * polter return Eq(diff(Gl, t), eq).subs({sympy.sympify('J_2(t)'): pi_g, sympy.sympify('J_1(t)'): q_g}) def G_eq(L): # massless neutrinos assert (L > 1) Gl = _make_index_func('G', L) Glp = _make_index_func('G', L + 1) Glm = _make_index_func('G', L - 1) eq = -k / (2 * L + 1) * ((L + 1) * Kf[L] * Glp - L * Glm) if L == 2: eq = eq + 8 * k / 15 * sigma return Eq(diff(Gl, t), eq).subs({sympy.sympify('G_2(t)'): pi_r, sympy.sympify('G_1(t)'): q_r}) def E_eq(L): # E polarization assert (L > 1) El = _make_index_func('E', L) Elp = _make_index_func('E', L + 1) Elm = _make_index_func('E', L - 1) eq = -k / (2 * L + 1) * ((L + 3) * (L - 1) * Kf[L] * Elp / (L + 1) - L * Elm) - opacity * El if L == 2: eq = eq + polter * opacity return Eq(diff(El, t), eq).subs(sympy.sympify('E_1(t)'), 0)
[docs] def get_hierarchies(lmax=5): """ Get Boltzmann hierarchies up to lmax for photons (J), E polarization and massless neutrinos (G). :param lmax: maxmimum multipole :return: list of equations """ eqs = [] for ell in range(2, lmax): eqs += [J_eq(ell), G_eq(ell), E_eq(ell)] return eqs
hierarchies = get_hierarchies() scalar_E_source = visibility * 15 * polter / 8 / (f_K(tau0 - t) * k) ** 2
[docs] def get_scalar_temperature_sources(checks=False): """ Derives terms in line of sight source, after integration by parts so that only integrated against a Bessel function (no derivatives). :param checks: True to do consistency checks on result :return: monopole_source, ISW, doppler, quadrupole_source """ # Line of sight temperature anisotropy source, by integrating by parts so only integrated against j_l source1 = k * sigma * exptau + visibility * 15 * polter / 8 source = source1 / 3 + diff(source1, t, t) / k ** 2 + visibility * Delta_g / 4 \ - hdot * exptau + diff(visibility * v_b - k * exptau * A, t) / k src = subs(var_subs, subs(background_eqs + pert_eqs, subs(background_eqs + pert_eqs, source).expand().doit()).simplify()).simplify().expand() ISW = src.coeff(exptau) * exptau if checks: # Check ISW is just 2*d phi/d eta assert (subs(Friedmann_Kfac_subs, subs(Friedmann_Kfac_subs, subs(var_subs, 2 * dphi * exptau - ISW).simplify()).simplify()).simplify() == 0) doppler1 = subs(var_subs, subs(pert_eqs, diff(visibility * (v_b + sigma), t) / k)).simplify().collect(visibility) remainder = src - ISW - doppler1 remainder = remainder.simplify().collect(visibility).subs(delta, solve(phi_sub, delta)).simplify() ISW = 2 * diff(phi, t) * exptau monopole_source = (Delta_g / 4 + (2 * phi + eta / (2 * K_fac))) * visibility quadrupole_source = (remainder - monopole_source).simplify() doppler = diff((v_b + sigma) * visibility, t) / k if checks: assert (subs(var_subs, subs(pert_eqs, (doppler - doppler1))).simplify() == 0) return monopole_source, ISW, doppler, quadrupole_source
# for translation in source output routine (does not include all equation variables) _camb_cache = {}
[docs] def camb_fortran(expr, name='camb_function', frame='CDM', expand=False): """ Convert symbolic expression to CAMB fortran code, using CAMB variable notation. This is not completely general, but it will handle conversion of Newtonian gauge variables like Psi_N, and most derivatives up to second order. :param expr: symbolic sympy expression using camb.symbolic variables and functions (plus any standard general functions that CAMB can convert to fortran). :param name: lhs variable string to assign result to :param frame: frame in which to interret non gauge-invariant expressions. By default uses CDM frame (synchronous gauge), as used natively by CAMB. :param expand: do a sympy expand before generating code :return: fortran code snippet """ camb_diff_vars = 'etakdot qgdot qrdot vbdot pigdot pirdot pinudot ' + \ 'octg octgdot polterdot polterddot diff_rhopi sigmadot phidot ' + \ 'ddvisibility dvisibility dopacity ddopacity' camb_arr_vars = 'Edot E' tau = _camb_cache.setdefault('tau', Symbol('tau')) etakdot, qgdot, qrdot, vbdot, pigdot, pirdot, pinudot, octg, octgdot, \ polterdot, polterddot, diff_rhopi, sigmadot, phidot, \ ddvisibility, dvisibility, dopacity, ddopacity = \ define_variables(camb_diff_vars, _camb_cache) Edot, E = sympy.symbols(camb_arr_vars, cls=sympy.IndexedBase, shape=(sympy.oo,)) # Keep everything except baryons pressure which is very small camb_diff_subs = [(diff(q_g, t), qgdot), (diff(pi_g, t), pigdot), (diff(pi_r, t), pirdot), (diff(q_r, t), qrdot), (diff(v_b, t), vbdot), (diff(visibility, t, t), ddvisibility), (diff(visibility, t), dvisibility), (diff(opacity, t, t), ddopacity), (diff(opacity, t), dopacity), (diff(J_3, t), octgdot), (diff(E_2, t), Edot[2]), (diff(E_3, t), Edot[3]), (diff(pi_nu, t), pinudot), (diff(eta, t), -2 * etakdot / k), (diff(Pi, t), diff_rhopi / kappa / a ** 2), (diff(polter, t, t), polterddot), (diff(polter, t), polterdot), (diff(sigma, t), sigmadot), (diff(phi, t), phidot), (diff(p_b, t), 0)] if frame != 'CDM': expr = make_frame_invariant(expr, frame) # substitute for variables not available in CAMB function expr = cdm_gauge(subs(Newt_vars + synchronous_vars + [hdot_sub, z_sub], expr)).doit().simplify() res = cdm_gauge(subs(background_eqs + total_eqs, expr).subs(camb_diff_subs)) if 'Derivative' in str(res): raise Exception( 'Unknown derivatives, generally can only handle up to second.\nRemaining derivatives: ' + str(res)) res = cdm_gauge(subs([K_sub, hdot_sub, z_sub], res)) camb_var_subs = [] for var in res.atoms(Function): camb_var = getattr(var, 'camb_var', None) if camb_var: if isinstance(camb_var, (list, tuple)): for x in camb_var: define_variable(x) camb_sub = getattr(var, 'camb_sub', None) if not camb_sub: raise Exception('must have camb_sub if camb_var has more than one variable') else: camb_var = define_variable(camb_var) camb_sub = getattr(var, 'camb_sub', None) or camb_var if camb_sub: if isinstance(camb_sub, str): camb_sub = eval(camb_sub) camb_var_subs.append((var, camb_sub)) camb_subs = camb_var_subs + [(p_b, 0), (E_2, E[2]), (E_3, E[3]), (J_3, octg), (K_fac, Kf[1])] res = res.subs(camb_subs).simplify() no_arg_funcs = [f for f in res.atoms(Function) if f.args[0] == t and f is not f_K] res = res.subs(zip(no_arg_funcs, [Symbol(str(x.func)) for x in no_arg_funcs])) res = res.subs(t, tau) if expand: res = res.expand() res = res.collect([Symbol(str(x.func)) for x in [k, sigma, opacity, visibility, dopacity, dvisibility, ddvisibility]]) res = sympy.fcode(res, source_format='free', standard=95, assign_to=name, contract=False) import textwrap if 'if ' not in res: lines = res.split('\n') for i, line in enumerate(lines): if '=' in line: res = '\n'.join(lines[i:]) break res = ''.join([x.strip() for x in res.split('&')]) res = ' &\n '.join(textwrap.wrap(res)) return res
_func_cache = {} _source_file_count = 0 _default_compiler = None _default_flags = None def get_default_compiler(): global _default_compiler, _default_flags if _default_compiler: return _default_compiler from .baseconfig import gfortran if gfortran: _default_compiler = 'gfortran' _default_flags = "-shared -fPIC -O1 -fmax-errors=4" else: import platform _default_compiler = 'ifort' if platform.system() == 'Darwin': _default_flags = "-dynamiclib -O1 -W0 -WB" else: _default_flags = "-shared -fpic -O1 -W0 -WB" # _default_flags="-shared -fPIC -g -fbounds-check -fbacktrace -ffpe-trap=invalid,overflow,zero", return _default_compiler _first_compile = True
[docs] def compile_source_function_code(code_body, file_path='', compiler=None, fflags=None, cache=True): """ Compile fortran code into function pointer in compiled shared library. The function is not intended to be called from python, but for passing back to compiled CAMB. :param code_body: fortran code to do calculation and assign sources(i) output array. Can start with declarations of temporary variables if needed. :param file_path: optional output path for generated f90 code :param compiler: compiler, usually on path :param fflags: options for compiler :param cache: whether to cache the result :return: function pointer for compiled code """ if cache and code_body in _func_cache: return _func_cache[code_body].source_func_ global _source_file_count template = """ REAL*8 function source_func(sources, tau, a, adotoa, grho, gpres,w_lam, cs2_lam, & grhob_t,grhor_t,grhoc_t,grhog_t,grhov_t,grhonu_t, & k,etak, etakdot, phi, phidot, sigma, sigmadot, & dgrho, clxg,clxb,clxc,clxr,clxnu, clxde, delta_p_b, & dgq, qg, qr, qde, vb, qgdot, qrdot, vbdot, & dgpi, pig, pir, pigdot, pirdot, diff_rhopi, & polter, polterdot, polterddot, octg, octgdot, E, Edot, & opacity, dopacity, ddopacity, visibility, dvisibility, ddvisibility, exptau, & tau0, tau_maxvis, Kf, f_K) implicit none real*8, intent(out) :: sources(:) REAL*8, intent(in) :: tau, a, adotoa, grho, gpres,w_lam, cs2_lam, & grhob_t,grhor_t,grhoc_t,grhog_t,grhov_t,grhonu_t, & k, etak, etakdot, phi, phidot, sigma, sigmadot, & dgrho, clxg,clxb,clxc,clxr,clxnu, clxde, delta_p_b, & dgq, qg, qr, qde, vb, qgdot, qrdot, vbdot, & dgpi, pig, pir, pigdot, pirdot, diff_rhopi, & polter, polterdot, polterddot, octg, octgdot, E(2:3), Edot(2:3), & opacity, dopacity, ddopacity, visibility, dvisibility, ddvisibility, exptau, & tau0, tau_maxvis, Kf(*) real*8, external :: f_K %s end function """ import subprocess import tempfile from ._compilers import is_32_bit, is_windows, compiler_environ, check_gfortran compiler = compiler or get_default_compiler() fflags = fflags or _default_flags if is_32_bit: fflags = "-m32 " + fflags if is_windows: global _first_compile fflags += ' -static' if _first_compile: check_gfortran(msg=True) workdir = file_path or tempfile.gettempdir() if not os.access(workdir, os.F_OK): os.mkdir(workdir) oldwork = os.getcwd() source_file = None try: os.chdir(workdir) _source_file_count += 1 while True: name_tag = 'camb_source%s' % _source_file_count dll_name = name_tag + '.dll' if not os.path.exists(dll_name): break try: os.remove(dll_name) except: _source_file_count += 1 source_file = name_tag + '.f90' with open(source_file, 'w') as f: f.write(template % code_body) command = " ".join([compiler, fflags, source_file, "-o", dll_name]) try: subprocess.check_output(command, stderr=subprocess.STDOUT, shell=True, cwd=workdir, env=compiler_environ) except subprocess.CalledProcessError as E: print(command) print('Error compiling generated code:') print(E.output) print('Source is:\n %s' % code_body) raise finally: if not file_path and source_file: os.remove(source_file) os.chdir(oldwork) # Had weird crashes when LoadLibrary path was relative to current dir dll_name = os.path.join(workdir, dll_name) func_lib = ctypes.LibraryLoader(ctypes.CDLL).LoadLibrary(dll_name) if cache: _func_cache[code_body] = func_lib if not file_path: # won't work on Windows while DLL in use try: os.remove(dll_name) except: pass return func_lib.source_func_
def compile_sympy_to_camb_source_func(sources, code_path=None, frame='CDM'): code = '' if not isinstance(sources, (list, tuple)): sources = [sources] for i, source in enumerate(sources): code += camb_fortran(source, 'sources(%s)' % (i + 1), frame=frame) + '\n' return compile_source_function_code(code, file_path=code_path) def internal_consistency_checks(): print('Sympy: ', sympy.__version__) # All equations should be gauge invariant for cons in constraints: assert (simplify(subs(Friedmann_Kfac_subs, make_frame_invariant(cons)) - cons) == 0) # Check deta equations consistent assert (subs(K_sub, subs(var_subs, subs(q_sub, deta).simplify() - subs(deta_2, diff(eta, t)))).simplify() == 0) # check consistency of constraint equations with evolution equations assert (subs(var_subs, subs(tot_eqs, diff(cons3, t))).simplify() == 0) assert (subs(Friedmann, subs(var_subs, subs(tot_eqs, diff(cons2, t))).simplify()).simplify() == 0) assert (subs(K_sub, subs(Friedmann, subs(var_subs, subs(tot_eqs, diff(cons1, t))).simplify())).simplify() == 0) # Check Weyl potential from Newtonian gauge vars assert (subs(var_subs, subs(pert_eqs, subs(Newt_vars, Phi_N + Psi_N - 2 * phi))).simplify() == 0) # Gauge functions checks assert (newtonian_gauge(dsigma) == 0) assert (subs(K_sub, subs(var_subs, subs(pert_eqs, subs(cdm_gauge(synchronous_vars), subs(synchronous_subs, sigma) - sigma).doit())) .simplify().collect(eta)).simplify() == 0) # check all equations are gauge invariant for eq in delta_eqs + vel_eqs: assert (simplify(subs(component_eqs + background_eqs, make_frame_invariant(eq.lhs).doit() - make_frame_invariant(eq.rhs))) == 0) try: # Check consistency of fluid equations with equations from total stress-energy conservation for eq in total_eqs: assert (subs(component_eqs + tot_subs, subs(tot_pert_subs, eq).doit()).doit().simplify()) except TypeError: print('If this test fails, you probably have an old sympy version. You need version 1 or higher.') raise get_scalar_temperature_sources(True) print("All symbolic relation tests OK") if __name__ == "__main__": internal_consistency_checks()