-
Notifications
You must be signed in to change notification settings - Fork 0
/
sampler.py
140 lines (120 loc) · 4.37 KB
/
sampler.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
from abc import ABC, abstractmethod
import numpy as np
from jetmontecarlo.utils.montecarlo_utils import *
########################################
# Base Sampler Class:
########################################
class sampler(ABC):
"""An abstract class designed to sample over phase space
for Monte Carlo integration
"""
# ------------------
# Sampling Methods:
# ------------------
@abstractmethod
def setArea(self):
"""Sets the area of the phase space associated
with the sampler."""
pass
@abstractmethod
def addSample(self):
"""Adds a single sample to the list self.samples,
and adds the associated jacobian factor of integration
to the list self.jacobians.
Not implemented in the abstract base class."""
pass
def generateSamples(self, num):
"""Generates some number of samples for the sampler."""
for _ in range(num):
self.addSample()
self.samples = np.array(self.samples)
def getSamples(self):
"""Returns the list of samples."""
return self.samples
def getSampleMethod(self):
"""Returns the sample method."""
return self.sampleMethod
def clearSamples(self):
"""Clears the list of samples."""
self.samples = []
# ------------------
# Saving/Loading:
# ------------------
def saveSamples(self, filename):
"""Saves the samples and area to filename.npz"""
np.savez(filename, samples=self.samples, area=self.area)
def loadSamples(self, filename):
"""Loads samples and area from filename.npz"""
npzfile = np.load(filename, allow_pickle=True,
mmap_mode='c')
self.samples = npzfile['samples']
self.area = npzfile['area']
# ------------------
# Validity Checks:
# ------------------
def checkSampleMethod(self):
"""Checking that the monte carlo sample method is
one of our supported sample methods."""
validSampleMethods = ['lin', 'log']
assert self.sampleMethod in validSampleMethods, \
str(self.sampleMethod) + "is not a supported sample method."
def checkValidEpsilon(self):
"""Checking that the keyword argument 'epsilon' is
valid, if the sampling method is logarithmic."""
eps = self.epsilon
eps = 0 if eps is None else eps
valid_eps = 0 < eps < 1.
if self.sampleMethod == 'log':
assert valid_eps, \
("epsilon={epsilon} is invalid for {type} sampling."
.format(epsilon=eps, type=self.sampleMethod))
@abstractmethod
# ------------------
# Init:
# ------------------
def __init__(self, sampleMethod, epsilon=0., **kwargs):
"""Prepares a MC sampler with
* The sample method ('lin'/'log') for the sample phase space, and
* any additional valid keyword arguments:
epsilon -- the lower cutoff of the logarithmic sampling.
"""
# Setting up:
self.sampleMethod = sampleMethod
self.epsilon = epsilon
# Initializing samples:
self.samples = []
self.jacobians = []
# Checking validity
self.checkSampleMethod()
self.checkValidEpsilon()
self.setArea()
# Working cooperatively with multiple inheritance:
super().__init__(**kwargs)
# ------------------------------------
# Simple sampler class:
# ------------------------------------
class simpleSampler(sampler):
"""A simple sampler which samples by default from 0 to 1."""
# ------------------
# Sampling:
# ------------------
def setArea(self):
if self.sampleMethod=='lin':
self.area = self.bounds[1] - self.bounds[0]
elif self.sampleMethod=='log':
self.area = np.log(1./self.epsilon)
def addSample(self):
if self.sampleMethod=='lin':
sample = getLinSample(self.bounds[0], self.bounds[1])
jac = 1.
if self.sampleMethod=='log':
sample = getLogSample(self.bounds[0], self.bounds[1], self.epsilon)
jac = sample
self.samples.append(sample)
self.jacobians.append(jac)
# ------------------
# Init:
# ------------------
def __init__(self, sampleMethod, bounds=[0,1], **kwargs):
self.bounds = bounds
super().__init__(sampleMethod, **kwargs)