-
Notifications
You must be signed in to change notification settings - Fork 7
/
Copy pathvector_fields.py
146 lines (112 loc) · 4.32 KB
/
vector_fields.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
import numpy as np
from numpy import linalg as la
import matplotlib.pyplot as plt
from custom_linalg import rotation_matrix_by_angle
class BaseField:
max_shift = .01
def direction(self, p):
''' returns unit vector '''
v = self._vector(p)
return v / (la.norm(v) + .000001)
def vector(self, p):
return self._vector(p)
def _vector(self, p):
return 1, 0
def point_close_to_source(self, p):
return False
def compute_traces(self, starting_points, trace_segments=100, dv=.02):
''' main logic here '''
num_traces = len(starting_points)
traces = np.zeros((num_traces, trace_segments, 2))
traces[:,0,:] = starting_points
for t in range(trace_segments - 1):
for i in range(num_traces):
if self.point_close_to_source(traces[i, t]):
traces[i, t + 1] = traces[i, t]
continue
delta = dv * self.vector(traces[i, t])
if la.norm(delta) > self.max_shift:
delta = self.direction(traces[i, t]) * self.max_shift
traces[i, t + 1] = traces[i, t] + delta
return traces
class ExpandingField(BaseField):
def _vector(self, p):
return p
class CosineField(BaseField):
def _vector(self, p):
x, y = p
v = np.array((np.cos(3 * (x + 2 * y)), np.sin(3 * (x - 2 * y))))
return 100 * v
class TwoCuspsField(BaseField):
def _vector(self, p):
x, y = p
v = np.array((x**2 + 2 * x * y, y**2 + 2 * x * y))
return 100 * v
class DipoleField(BaseField):
''' someone on Internet said this is expression for dipole field '''
def _vector(self, p):
x, y = 10 * p
v = np.array(((x + 1) / ((y + 1)**2 + y**2) - (x - 1) / ((x - 1)**2 + y**2),
y / ((y + 1)**2 + x**2) - y / ((x - 1)**2 + y**2)
))
return 100 * v
class CurlField(BaseField):
''' CurlField is a compostion of spiral fields produced by sources.
`sources` is a tuple of tuples with coordinates of some 'particles'
and direction of spiral (in radians) relative to source position '''
sources = ((np.array((.1, .2)), np.pi/2),
(np.array((.1, .9)), np.pi/6),
(np.array((.7, .9)), 0))
def forces(self, p):
return [( (p - src) @ rotation_matrix_by_angle(angle) )
/ (la.norm(p - src) ** 2 + 0.0001)
for src, angle in self.sources]
def __init__(self, sources=None):
if sources:
self.sources = sources
def _vector(self, p):
return sum(self.forces(p), np.array([0, 0]))
def point_close_to_source(self, p):
for src, _ in self.sources:
if la.norm(src - p) < .005:
return True
return False
class DivCurlField(BaseField):
''' this was initial version of CurlField '''
sources = ((np.array((.6, .2)), 'curl', 1),
(np.array((.2, .9)), 'div', .5))
def forces(self, p):
rotation = {
'curl': np.array([[0, -1], [1, 0]]),
'div': np.identity(2),
}
return [( (p - src) @ rotation[_type] * mass )
/ ( la.norm(p - src)**2 + .001 )
for src, _type, mass in self.sources]
def __init__(self, sources=None):
if sources:
self.sources = sources
def _vector(self, p):
return sum(self.forces(p), np.array([0, 0]))
fields = (ExpandingField,
CosineField,
TwoCuspsField,
DipoleField,
DivCurlField,
CurlField,
)
def preview_flow(field, n_traces=100, trace_segments=15,
dv=.01, dots=False, starting_points=None, subplot=None):
if not subplot:
_, subplot = plt.subplots()
setup_empty_subplot(subplot, field.__class__.__name__)
if not starting_points:
starting_points = np.random.rand(n_traces, 2) - np.array((.5, .5))
traces = field.compute_traces(starting_points, trace_segments, dv=dv)
for trace in traces:
subplot.plot(*trace.T, color='grey')
if dots:
subplot.scatter(*traces[:,0,:].T, color='black', s=3)
def setup_empty_subplot(subplot, title=None):
subplot.axis('equal')
subplot.set_title(title)