Collisional-Rings/botsrots.py at main · rieder/Collisional-Rings · 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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
import sys
import numpy

from amuse.units import units
from amuse.datamodel import Particles
from amuse.support.exceptions import AmuseException

from amuse.units.quantities import VectorQuantity
import numpy as np


class BotsRots(object):
    """
    Resolves collisions between particles by treating them as somewhat inelastic
    collisions. The degree of rebounding is controlled by the epsilon_n parameter.

    """

    stellar_evolution_code_required = False
    gravity_code_required = False

    def __init__(self, epsilon_n=0.1):
        if 0 <= epsilon_n <= 1:
            self.epsilon_n = epsilon_n
        else:
            raise AmuseException("epsilon_n must be in the range [0, 1]")

    def handle_collisions(self, particles, primary_indices, secondary_indices):
        d_pos = VectorQuantity(
                np.zeros(len(particles)*3).reshape(len(particles),3),
                particles.position.unit,
                )
        d_vel = VectorQuantity(
                np.zeros(len(particles)*3).reshape(len(particles),3),
                particles.velocity.unit,
                )

        # This loop could be run with many threads...
        primaries   = particles[primary_indices]
        secondaries = particles[secondary_indices]

        relative_position   = secondaries.position - primaries.position
        distance            = relative_position.lengths()
        relative_velocity   = secondaries.velocity - primaries.velocity
        surface_distance    = distance - primaries.radius - secondaries.radius
        total_mass          = primaries.mass + secondaries.mass

        normal  = relative_position / distance.reshape((len(primaries),1))
        nx      = normal[:,0]
        ny      = normal[:,1]
        nz      = normal[:,2]
        normal_velocity = (
                relative_velocity[:,0]*nx +
                relative_velocity[:,1]*ny +
                relative_velocity[:,2]*nz
                )

        pshift = VectorQuantity(
                np.zeros(len(primaries)*3).reshape(len(primaries),3),
                (particles.position.unit / particles.mass.unit),
                )
        vshift = VectorQuantity(
                np.zeros(len(primaries)*3).reshape(len(primaries),3),
                (particles.velocity.unit / particles.mass.unit),
                )
        pshift[:,0] = (1./total_mass * surface_distance * nx)
        pshift[:,1] = (1./total_mass * surface_distance * ny)
        pshift[:,2] = (1./total_mass * surface_distance * nz)
        vshift[:,0] = (1./total_mass * (1.0+self.epsilon_n) * normal_velocity * nx)
        vshift[:,1] = (1./total_mass * (1.0+self.epsilon_n) * normal_velocity * ny)
        vshift[:,2] = (1./total_mass * (1.0+self.epsilon_n) * normal_velocity * nz)
        dpos_pri = secondaries.mass.reshape((len(primaries),1)) * pshift
        dvel_pri = secondaries.mass.reshape((len(primaries),1)) * vshift
        dpos_sec = -primaries.mass.reshape((len(primaries),1)) * pshift
        dvel_sec = -primaries.mass.reshape((len(primaries),1)) * vshift
        for i in range(len(primary_indices)):
            d_pos[primary_indices[i]] += dpos_pri[i]
            d_vel[primary_indices[i]] += dvel_pri[i]
            d_pos[secondary_indices[i]] += dpos_sec[i]
            d_vel[secondary_indices[i]] += dvel_sec[i]

        return d_pos, d_vel