-
Notifications
You must be signed in to change notification settings - Fork 75
/
tree.py
49 lines (36 loc) · 1.56 KB
/
tree.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
"""Gravity Tree. Builds tree based on pkdgrav2"""
try:
from .. import pkdgrav
except ImportError:
# pkdgrav2 never works at the moment, and the warning below
# confuses/annoys people
pass
#import warnings
#warnings.warn("Unable to import PKDGrav gravity solver. Most likely this means either that your installation is broken, or that you are running python inside the pynbody distribution directory, in which case python cannot see the installed version. However, it also doesn't matter unless you want to use the tree gravity solver.",RuntimeWarning)
from time import process_time
import numpy as np
from .. import config
class GravTree:
def __init__(self, pos, mass, rs, eps=None, leafsize=16):
start = process_time()
self.tree = pkdgrav.pkdPythonInitialize(
pos, mass, eps, rs, int(leafsize))
end = process_time()
if config['verbose']:
print('Tree build done in %5.3g s' % (end - start))
self.derived = True
self.flags = {'WRITEABLE': False}
def calc(self, vec_pos, eps=None):
accel = np.zeros((len(vec_pos), 3))
pot = np.zeros(len(vec_pos))
if config['verbose']:
print('Calculating Gravity')
start = process_time()
pkdgrav.pkdPythonDoGravity(self.tree, accel, pot, theta=0.55)
end = process_time()
if config['verbose']:
print('Gravity calculated in %5.3g s' % (end - start))
return accel, pot
def __del__(self):
if hasattr(self, 'tree'):
pkdgrav.free(self.tree)