forked from ibackus/ICgen
-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathmake_sigma.py
239 lines (175 loc) · 7.36 KB
/
make_sigma.py
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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
# -*- coding: utf-8 -*-
"""
Calculates cubic spline interpolations for sigma(r) and probability(r)
probability = 2*pi*r*sigma
Created on Mon Jan 27 13:00:52 2014
@author: ibackus
"""
# ICgen packages
import isaac
# External packages
import pynbody
SimArray = pynbody.array.SimArray
import numpy as np
from scipy.interpolate import interp1d
from scipy.integrate import quad
from scipy.integrate import simps
import copy as copier
class sigma_gen:
"""
A class to generate the surface density (sigma), probability density (pdf)
and inverse cumulative distribution function (cdf_inv) as a function of r
USAGE:
# Generate sigma (settings is generated by ICgen_settings.py, see below)
import make_sigma
sigma = make_sigma.sigma_gen(r, sigma)
# Calculate at various r positions:
sigma(r) # returns sigma evaluated at r
sigma.sigma(r) # returns sigma evaluated at r
pdf = sigma.pdf(r) # returns pdf evaluated at r
cdf_inv = sigma.cdf_inv(m) # returns cdv_inv at m for 0 < m < 1
# Generate sigma with a precalulated CDF (as done before)
sigma = make_sigma.sigma_gen(r, sigma, CDF)
"""
def __init__(self, r_bins, sigmaBinned, CDF=None):
self.input_dict = {'r': r_bins, 'sigma': sigmaBinned}
self._make_sigma(r_bins, sigmaBinned)
self._make_pdf()
self._make_cdf_inv(CDF)
self._disk_mass()
def __call__(self, r):
return self.sigma(r)
def _make_sigma(self, r_bins, sigmaBinned):
"""
Generates the surface density as a function of r, a callable object
sigma(r) and assigns it to self.sigma
Generates a spline interpolation of sigma vs r from the file
defined by settings.filenames.sigmaFileName. Returns sigma vs r as
an cubic spline interpolation object
(see scipy.interpolation.interp1d)
sigma_input should be a pickled dictionary with the entries:
'sigma': <sigma evaluated at r>
'r': <r for the bins>
If the input sigma has units, sigma vs r will be returned in units
of Msol/au^2
"""
# Convert to default units of Msol/au^2. If no units, assign default
sigmaBinned = isaac.match_units(sigmaBinned, 'Msol au**-2')[0]
# Convert r_bins to default units of 'au'
r_bins = isaac.match_units(r_bins, 'au')[0]
# Calculate spline interpolation
print 'Calculating spline interpolation (slow for many data points)'
sigspline = interp1d(r_bins,sigmaBinned,kind='cubic',fill_value=0.0,\
bounds_error=False)
def sigout(r):
"""
Linear spline interpolation of sigma(r).
ARGUMENTS:
r - can be scalar, numpy array, or sim array
RETURNS:
sigma (surface density) evaluated at r
"""
# Try to convert r to the units used to make sigspline ('au')
r = isaac.match_units(r, 'au')[0]
return SimArray(sigspline(r), 'Msol au**-2')
self.sigma = sigout
self.r_bins = r_bins
def _make_pdf(self):
"""
Generates the probability density as a function of r (from the surface
density) and returns it as a callable function pdf(r) to self.pdf
pdf(r) = 2*pi*r*sigma(r), up to a normalization
Using sigma(r) calculated in _make_sigma and r_bins loaded from
settings.filenames.sigmaFileName, creates an interpolation over
r_bins, sigma(r_bins). The PDF will be approximately normalized
"""
# Calculate the binned PDF
pdfBinned = 2*np.pi*self.r_bins * self.sigma(self.r_bins)
# Normalize
integral = simps(pdfBinned,self.r_bins)
pdfBinned /= integral
# Calculate a spline interpolation
print 'Calculating spline interpolation (slow for many data points)'
pdfSpline = interp1d(self.r_bins, pdfBinned, kind='cubic',\
fill_value=0.0, bounds_error=False)
def pdf_fcn(r_in):
"""
Normalized cubic spline interpolation of the PDF(r) from sigma(r).
The PDF is just calculated as 2*pi*r*sigma(r).
ARGUMENTS:
r_in - radii at which to calculate the PDF
RETURNS:
probability density function from sigma(r) evaluated at r_in
"""
# Put r_in into the units used in generating the pdf
r_in = isaac.match_units(r_in, self.r_bins)[0]
# Evaluate the pdf at r_in
pdf_vals = pdfSpline(r_in)
# Put the pdf into units of r_in.units**-1
pdf_vals = isaac.match_units(pdf_vals, 1/r_in)[0]
return pdf_vals
self.pdf = pdf_fcn
def _make_cdf_inv(self, f=None):
"""
Calculates the inverse of the CDF (cumulative distribution function).
This can be use for generating particle positions. Generates a
callable method and saves to self.cdf_inv
The CDF_inv is made by cumulatively integrating the PDF over the radial
bins defined in self.r_bins
The optional argument, f, is the CDF binned over the radial bins
"""
print 'calculating CDF'
# Calculate the CDF from prob
r = self.r_bins
r[0] = 0.0
nr = len(r)
if f is None:
f = np.zeros(nr)
for n in range(nr):
f[n] = quad(self.pdf,r[[0]],r[[n]])[0]
f /= f.max()
self._CDF = f.copy()
print 'calculating inverse CDF'
# Calculate the inverse CDF.
# Assume CDF is approximately monotonic and sort to force it to be
ind = f.argsort()
f = f[ind]
r = r[ind]
# Drop values where CDF is constant (ie, prob = 0)
mask = np.ones(nr,dtype='bool')
for n in range(1,nr):
if f[n] == f[n-1]:
mask[n] = False
f = f[mask]
r = r[mask]
finv = interp1d(f,r,kind='linear')
def finv_fcn(m_in):
"""
The inverse CDF for sigma(r).
input:
0 <= m_in < 1
returns:
r (radius), the inverse CDF evaluated at m_in
Uses a linear spline interpolation.
"""
r_out = finv(m_in)
r_out = isaac.match_units(r_out, r)[0]
return r_out
self.cdf_inv = finv_fcn
def _disk_mass(self):
"""
Calculate the total disk mass by integrating sigma
"""
# Assign variables
r = self.r_bins
sig = self.sigma(r)
# Now integrate
m_disk = simps(2*np.pi*r*sig, r)
m_units = sig.units * (r.units)**2
m_disk = isaac.match_units(m_disk, m_units)[0]
self.m_disk = m_disk
def copy(self):
"""
Returns a copy of the sigma object
"""
return copier.copy(self)