general_purpose_python_scripts/ThreeDimensionalFunction.py~ at main · stubbslab/general_purpose_python_scripts · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
import numpy as np
import math
from scipy.interpolate import RegularGridInterpolator
from scipy.interpolate import interp1d
import scipy.integrate as integrate

class ThreeDimensionalFunction:

    def integrate_los(self, xs = None, zs = None, integration_bin_centers = None, los_bin_step = 1.0): #,integration_rect_centers):
        if xs is None: xs = self.xs
        if zs is None: zs = self.zs
        print 'xs = ' + str(xs)
        print 'zs = ' + str(zs)
        if integration_bin_centers is None: integration_bin_centers = self.bin_centers
        n_centers = len(integration_bin_centers)
        integration_bin_borders = ( [integration_bin_centers[0] - (integration_bin_centers[1] - integration_bin_centers[0]) / 2.0]
                                        + [(integration_bin_centers[i+1] + integration_bin_centers[i]) / 2.0 for i in range(n_centers - 1)]
                                        + [integration_bin_centers[n_centers-1] + (integration_bin_centers[n_centers-1] - integration_bin_centers[n_centers - 2]) /2.0 ] )
        integration_bin_borders.sort()
        integrated_array=np.zeros((len(xs),len(zs)))

        integrated_array = sum([ self.getRotatedCrossSection(self.funct, xs, integration_bin_centers[i], zs)
                                  * (integration_bin_borders[i + 1] - integration_bin_borders[i]) for i in range(len(integration_bin_centers))])
        #print 'integrated_array = ' + str(integrated_array)
        onSkyInterpolator = RegularGridInterpolator((xs, zs), integrated_array, method = 'linear')

        return onSkyInterpolator


    def getRotatedCrossSection(self, funct, xs, y_val, zs):
        xmesh, zmesh = np.meshgrid(xs, zs)
        return np.transpose(funct(xmesh, y_val, zmesh))

    def getRadialAverageMassFunction(self, Rs, phis, xs = None, zs = None, funct = None, integration_bin_centers = None):
        print 'len(Rs) = ' + str(len(Rs))
        if xs is None: xs = self.xs
        if zs is None: zs = self.zs
        if funct is None: funct = self.funct
        if integration_bin_centers is None: integration_bin_centers = self.bin_centers
        x_of_phi = lambda R, phi: R * math.cos(phi)
        z_of_phi = lambda R, phi: R * math.sin(phi)
        surface_mass_interpolator = self.integrate_los(xs = xs, zs = zs, integration_bin_centers = integration_bin_centers)
        print integrate.quad(lambda phi: surface_mass_interpolator( (x_of_phi((Rs[3] + Rs[2])/2.0, phi), z_of_phi((Rs[3] + Rs[2])/2.0, phi) )), 0, 2.0*math.pi )
        print [surface_mass_interpolator( (x_of_phi(Rs[1], phis[j]), z_of_phi(Rs[1], phis[j])) ) * (phis[j+1] - phis[j]) for j in range(len(phis) - 1) ]
        print 'Here 1'
        for i in range(len(Rs) - 1):
            print sum([surface_mass_interpolator( (x_of_phi((Rs[i+1] + Rs[i])/2.0, phis[j]), z_of_phi((Rs[i+1] + Rs[i])/2.0, phis[j])) )
                                       * (phis[j+1] - phis[j]) for j in range(len(phis) - 1) ])
        ring_integrated_masses = [(Rs[i+1] - Rs[i])/2.0 * (Rs[i+1] + Rs[i])/2.0 *
                                  sum([surface_mass_interpolator( (x_of_phi((Rs[i+1] + Rs[i])/2.0, phis[j]), z_of_phi((Rs[i+1] + Rs[i])/2.0, phis[j])) )
                                       * (phis[j+1] - phis[j]) for j in range(len(phis) - 1) ]) for i in range(len(Rs)-1)]
        ring_integrated_average_masses = [sum([surface_mass_interpolator( (x_of_phi((Rs[i+1] + Rs[i])/2.0, phis[j]), z_of_phi((Rs[i+1] + Rs[i])/2.0, phis[j])) )
                                          * (phis[j+1] - phis[j]) for j in range(len(phis) - 1) ]) for i in range(len(Rs)-1)]

        radial_interpolator = interp1d([(Rs[i+1] + Rs[i])/2.0 for i in range(len(Rs)-1)], ring_integrated_average_masses, kind = 'cubic')
        return radial_interpolator

    def __init__(self, funct, xs, bin_centers, zs):
        self.funct = funct
        self.xs = xs
        self.bin_centers = bin_centers
        self.zs = zs
        self.surfaceMassInterpolator = self.integrate_los(xs = self.xs, zs = self.zs, integration_bin_centers = self.bin_centers)