Source code for pyflation.sourceterm.sosource

"""sosource.py Second order source term calculation module.

Provides the method getsourceandintegrate which uses an instance of a first
order class from cosmomodels to calculate the source term required for second
order models.


"""
#Author: Ian Huston
#For license and copyright information see LICENSE.txt which was distributed with this file.



from __future__ import division # Get rid of integer division problems, i.e. 1/2=0

import numpy as np
import logging
import time
import os

from pyflation.sourceterm import srcequations 
from pyflation import configuration

#Set debug flag
_debug = configuration._debug


if not "profile" in __builtins__:
[docs] def profile(f): return f #Start logging
source_logger = logging.getLogger(__name__)
[docs]def set_log_name(): root_log_name = logging.getLogger().name source_logger.name = root_log_name + "." + __name__
[docs]def calculatesource(m, nix, integrand_elements, srceqns): """Return the integrated source term at this timestep. Given the first order model and the timestep calculate the integrated source term at this time step. Parameters ---------- m: Cosmomodels.FOCanonicalTwoStage instance First order model to be used. nix: int Index of current timestep in m.tresult integrand_elements: tuple Contains integrand arrays in order integrand_elements = (k, q, theta) srceqns: Instance of a subclass of srcequations.SourceEquations This class should implement the sourceterm function which constructs the integrated source term. Returns ------- s: array_like Integrated source term calculated using srcfunc. """ k, q, theta = integrand_elements #Copy of yresult myr = m.yresult[nix].copy() #Fill nans with initial conditions if np.any(m.tresult[nix] < m.fotstart): #Get first order ICs: nanfiller = m.getfoystart(m.tresult[nix].copy(), np.array([nix])) if _debug: source_logger.debug("Left getfoystart. Filling nans...") #switch nans for ICs in m.yresult are_nan = np.isnan(myr) myr[are_nan] = nanfiller[are_nan] if _debug: source_logger.debug("NaNs filled. Setting dynamical variables...") #Get first order results #Check that first order model is single field if m.nfields != 1: raise NotImplementedError("Source term can only be calculated for single field models.") #The background variables can just be taken as the first k value bgvars = myr[m.bg_ix,0] #The first order variables need the full k values #Until the source term calculation is upgraded to use multiple fields #just use the first (and only) field here dphi1 = myr[m.dps_ix][0] dphi1dot = myr[m.dpdots_ix][0] #Setup interpolation if _debug: source_logger.debug("Variables set. Getting potentials for this timestep...") #To get the potentials we send the recorded values to the potentials function arraypotentials = list(m.potentials(myr, m.pot_params)) #Get potentials in right shape #Shape of potentials is now determined by number of fields and not ks. #So no need to do this for each potential, but need to be careful about use #of potentials later on. #Warning: This should be changed when source term calculation can do multiple fields potentials = [p if np.isscalar(p) else np.asscalar(p) for p in arraypotentials] #If the potential is k dependent this needs to be changed to include full #k behaviour. The bgvars variable should also include the full k range in #this case. #Value of a for this time step a = m.ainit*np.exp(m.tresult[nix]) if _debug: source_logger.debug("Calculating source term integrand for this timestep...") #Get integrated source term src = srceqns.sourceterm(bgvars, a, potentials, dphi1, dphi1dot) if _debug: source_logger.debug("Integration successful!") return src
[docs]def getsourceandintegrate(m, savefile=None, srcclass=None, ninit=0, nfinal=-1, ntheta=513, numks=1025, overwrite=False): """Calculate and save integrated source term. Using first order results in the specified model, the source term for second order perturbations is obtained from the given source function. The convolution integral is performed and the results are saved in a file with the specified filename. Parameters ---------- m: compatible cosmomodels model instance The model class should contain first order results as in `cosmomodels.FOCanonicalTwoStage` savefile: String, optional Filename where results should be saved. srcclass: class, optional Class which contains functions for constructing source term. Defaults to srcequations.SlowRollSource. Class should be a subclass of srcequations.SourceEquations and should implement the sourceterm function ninit: int, optional Start time index for source calculation. Default is 0 (start at beginning). nfinal: int, optional End time index for source calculation. -1 signifies end of run (and is the default). ntheta: int, optional Number of theta points to integrate. Should be a power of two + 1 for Romberg integration. Default is 129. numks: int, optional Number of k modes required for second order calculation. Must be small enough that largest first order k mode is at least 2*(numks*deltak + kmin). Default is this limiting value. Returns ------- savefile: String Filename where results have been saved. """ #Check time limits if ninit < 0: ninit = 0 if nfinal > m.tresult.shape[0] or nfinal == -1: nfinal = m.tresult.shape[0] #Change to ints for xrange ninit = int(ninit) nfinal = int(nfinal) if not numks: numks = round((m.k[-1]/2 - m.k[0])/(m.k[1]-m.k[0])) + 1 #Initialize variables for all timesteps k = q = m.k[:numks] #Need np=len(m.k)/2 for case when q=-k #Check consistency of first order k range if (m.k[-1] - 2*k[-1])/m.k[-1] < -1e-12: raise ValueError("First order k range not sufficient!") theta = np.linspace(0, np.pi, ntheta) firstmodestart = min(m.fotstart) lastmodestart = max(m.fotstart) #Check source class and create instance if srcclass is None: srcclass = srcequations.SlowRollSource if not issubclass(srcclass, srcequations.SourceEquations): raise TypeError("Source class should be a class derived from srcequations.SourceEquations.") kfx = {"kmin":k[0], "deltak":k[1]-k[0], "numsoks":numks, "fullkmax":m.k[-1], "nthetas":ntheta} srceqns = srcclass(kfx) #Pack together in tuple integrand_elements = (k, q, theta) #Main try block for file IO try: sf, sarr, narr = opensourcefile(k, savefile, sourcetype="term", overwrite=overwrite) try: # Begin calculation if _debug: source_logger.debug("Entering main time loop...") #Main loop over each time step for nix in xrange(ninit, nfinal): if nix%10 == 0: source_logger.info("Starting nix=%d/%d sequence...", nix, nfinal) #Only run calculation if one of the modes has started. if m.tresult[nix] > lastmodestart or m.tresult[nix+2] >= firstmodestart: src = calculatesource(m, nix, integrand_elements, srceqns) else: src = np.nan*np.ones_like(k) sarr.append(src[np.newaxis,:]) narr.append(np.array([nix])) if _debug: source_logger.debug("Results for this timestep saved.") finally: #source = np.array(source) sf.close() except IOError: raise return savefile
[docs]def opensourcefile(k, filename=None, sourcetype=None, overwrite=False): """Open the source term hdf5 file with filename.""" import tables #Set up file for results if not filename or not os.path.isdir(os.path.dirname(filename)): source_logger.info("File or path to file %s does not exist." % filename) date = time.strftime("%Y%m%d%H%M%S") filename = os.path.join(os.getcwd(), "src" + date + ".hf5") source_logger.info("Saving source results in file " + filename) if not sourcetype: raise TypeError("Need to specify filename and type of source data to store [int(egrand)|(full)term]!") if sourcetype in ["int", "term"]: sarrname = "source" + sourcetype if _debug: source_logger.debug("Source array type: " + sarrname) else: raise TypeError("Incorrect source type specified!") #Check if file exists and set write flags depending on overwrite option if os.path.isfile(filename): if overwrite: source_logger.info("File %s exists and will be overwritten." % filename) writeflag = "w" else: source_logger.info("File %s exists and results will be appended." % filename) writeflag = "a" else: writeflag = "w" #Add compression to files and specify good chunkshape filters = tables.Filters(complevel=1, complib=configuration.hdf5complib) #cshape = (10,10,10) #good mix of t, k, q values #Get atom shape for earray atomshape = (0, len(k)) try: if _debug: source_logger.debug("Trying to open source file " + filename) rf = tables.openFile(filename, writeflag, "Source term result") if not "results" in rf.root: if _debug: source_logger.debug("Creating group 'results' in source file.") resgrp = rf.createGroup(rf.root, "results", "Results") else: resgrp = rf.root.results if not sarrname in resgrp: if _debug: source_logger.debug("Creating array '" + sarrname + "' in source file.") sarr = rf.createEArray(resgrp, sarrname, tables.ComplexAtom(itemsize=16), atomshape, filters=filters) karr = rf.createEArray(resgrp, "k", tables.Float64Atom(), (0,), filters=filters) narr = rf.createEArray(resgrp, "nix", tables.IntAtom(), (0,), filters=filters) karr.append(k) else: if _debug: source_logger.debug("Source file and node exist. Testing source node shape...") sarr = rf.getNode(resgrp, sarrname) narr = rf.getNode(resgrp, "nix") if sarr.shape[1:] != atomshape[1:]: raise ValueError("Source node on file is not correct shape!") except IOError: raise return rf, sarr, narr