def add_result (result_name, qp_v):
nonlocal results
H_v = vorpy.apply_along_axes(PendulumNd.H, (-2,-1), (qp_v,), output_axis_v=(), func_output_shape=())
deviation_form_v = deviation_form(t_v=t_v, qp_v=qp_v, dH_dq=PendulumNd.dH_dq, dH_dp=PendulumNd.dH_dp)
norm_deviation_form_v = vorpy.apply_along_axes(np.linalg.norm, (-2,-1), (deviation_form_v,), output_axis_v=(), func_output_shape=())
norm_error_v = np.full(norm_deviation_form_v.shape, np.nan)
results.add_result(result_name, dt, t_v, qp_v, H_v, norm_deviation_form_v, norm_error_v)
def test__apply_along_axes__compare_with__apply_along_axis():
rng = np.random.RandomState(42)
a = rng.randn(4, 5, 6, 7)
assert np.all(
vorpy.apply_along_axes(np.sum, [0], [a]) == np.apply_along_axis(
np.sum, 0, a))
assert np.all(
vorpy.apply_along_axes(np.sum, [1], [a]) == np.apply_along_axis(
np.sum, 1, a))
assert np.all(
vorpy.apply_along_axes(np.sum, [2], [a]) == np.apply_along_axis(
np.sum, 2, a))
assert np.all(
vorpy.apply_along_axes(np.sum, [3], [a]) == np.apply_along_axis(
np.sum, 3, a))
assert np.all(
vorpy.apply_along_axes(np.sum, (0, ), (
a, )) == np.apply_along_axis(np.sum, 0, a))
assert np.all(
vorpy.apply_along_axes(np.sum, (1, ), (
a, )) == np.apply_along_axis(np.sum, 1, a))
assert np.all(
vorpy.apply_along_axes(np.sum, (2, ), (
a, )) == np.apply_along_axis(np.sum, 2, a))
assert np.all(
vorpy.apply_along_axes(np.sum, (3, ), (
a, )) == np.apply_along_axis(np.sum, 3, a))
print('test__apply_along_axes__compare_with__apply_along_axis passed.')
def test__salvaged_result ():
import sys
def dV_dq_complaining (q):
# Introduce arbitrary error in order to test SalvagedResultException functionality.
assert np.sum(np.square(q)) > 1.0e-2
return KeplerNd.dV_dq(q)
dt = 0.01
t_v = np.arange(0.0, 60.0, dt)
# R^3 Kepler problem
qp_0 = np.array([
[1.0, 0.0, 0.0],
[0.0, 0.01, 0.0]
])
order = 2
omega = np.pi/(4*dt)
assert np.allclose(2*omega*dt, np.pi/2)
try:
results = Results()
qp_v = vorpy.symplectic_integration.separable_hamiltonian.integrate(
initial_coordinates=qp_0,
t_v=t_v,
dK_dp=KeplerNd.dK_dp,
dV_dq=dV_dq_complaining,
update_step_coefficients=vorpy.symplectic_integration.separable_hamiltonian.update_step_coefficients.ruth4
)
assert False, 'did not catch SalvagedResultException as expected'
except vorpy.symplectic_integration.exceptions.SalvagedResultException as e:
print('caught SalvagedResultException as expected: {0}'.format(e))
qp_v = e.salvaged_qp_v
t_v = e.salvaged_t_v
assert qp_v.shape[0] == t_v.shape[0]
H_v = vorpy.apply_along_axes(KeplerNd.H, (-2,-1), (qp_v,), output_axis_v=(), func_output_shape=())
deviation_form_v = deviation_form(t_v=t_v, qp_v=qp_v, dH_dq=KeplerNd.dH_dq, dH_dp=KeplerNd.dH_dp)
norm_deviation_form_v = vorpy.apply_along_axes(np.linalg.norm, (-2,-1), (deviation_form_v,), output_axis_v=(), func_output_shape=())
norm_error_v = np.full(norm_deviation_form_v.shape, np.nan)
results.add_result('Kepler trajectory', dt, t_v, qp_v, H_v, norm_deviation_form_v, norm_error_v)
# As you can see in this plot, the energy is not nearly conserved as time goes on, and the norm of the deviation form
# diverges away from zero as time goes on.
results.plot(filename=make_filename_in_artifacts_dir('symplectic_integration.separable_hamiltonian.salvaged_result.png'))
def test__apply_along_axes():
def symmetric_square(m):
return np.einsum('ij,kj->ik', m, m)
def is_symmetric(m):
return np.all(m == m.T)
rng = np.random.RandomState(42)
a = rng.randn(2, 3, 4, 5)
N = len(a.shape)
# Use all possible combinations of input and output axes.
for input_i0 in range(-N, N - 1):
for input_i1 in range(input_i0 + 1, N):
# Only test the pairs where input_i0 indexes an axis before that indexed by input_i1.
if input_i0 + N >= input_i1:
continue
for output_i0 in range(-N, N - 1):
for output_i1 in range(output_i0 + 1, N):
# Only test the pairs where output_i0 indexes an axis before that indexed by output_i1.
if output_i0 + N >= output_i1:
continue
output_axis_v = (output_i0, output_i1)
# Compute the result. The multi-slice across the output axes should be a symmetric matrix.
result = vorpy.apply_along_axes(
symmetric_square, (input_i0, input_i1), (a, ),
output_axis_v=output_axis_v)
# Figure out which indices correspond to the input axes; call these result_non_output_axis_v.
all_indices = tuple(range(N))
normalized_output_axis_v = tuple(
output_axis if output_axis >= 0 else output_axis + N
for output_axis in output_axis_v)
result_non_output_axis_v = sorted(
list(
frozenset(all_indices) -
frozenset(normalized_output_axis_v)))
# print('output_axis_v = {0}, result_non_output_axis_v = {1}'.format(output_axis_v, result_non_output_axis_v))
assert len(result_non_output_axis_v) == 2
# Take all multi-slices and verify symmetry.
for check_i0, check_i1 in itertools.product(
range(result.shape[result_non_output_axis_v[0]]),
range(result.shape[result_non_output_axis_v[1]])):
# Construct the multi-slice to take.
multislice = [slice(None) for _ in range(4)]
multislice[result_non_output_axis_v[0]] = check_i0
multislice[result_non_output_axis_v[1]] = check_i1
# print('multislice = {0}'.format(multislice))
assert is_symmetric(result[tuple(multislice)])
print('test__apply_along_axes passed.')
def deviation_form(*, t_v, qp_v, dH_dq, dH_dp):
"""
Hamilton's equations are
dq/dt = \partial H / \partial p
dp/dt = - \partial H / \partial q
This function computes the deviation of the discretized curve given by time values t_v and coordinates
qp_v from satisfaction of Hamilton's equations. A perfect solution would have a deviation of zero at all
time/point pairs (though because we're dealing with a discretized curve, this wouldn't imply that the
solution is exact).
The value returned is the evaluation of a form equivalent to Hamilton's equations:
dq/dt - \partial H / \partial p
dp/dt + \partial H / \partial q
and has the shape (T,A_1,A_2,...,A_K,2,N), which will be the same shape as qp_v, noting that K may be 0.
Note that because t_v indexes the time values corresponding to the elements of qp_v.
The [discrete] derivatives dq/dt and dp/dt are computed using a symmetric difference quotient:
(dq/dt)[i] := (q[i+1] - q[i-1]) / (t[i+1] - t[i-1])
with (dp/dt)[i] defined analogously. For the first and last entries, a non-symmetric difference quotient
will be used.
(dq/dt)[ 0] := (q[ 1] - q[ 0]) / (t[ 1] - t[ 0])
(dq/dt)[-1] := (q[-1] - q[-2]) / (t[-1] - t[-2])
with (dp/dt)[0] and (dp/dt)[-1] defined analogously.
TODO: Maybe only compute this for interior samples; end-point difference quotients might not be that useful.
"""
#print('deviation_form; qp_v.shape = {0}'.format(qp_v.shape))
assert np.all(
np.diff(t_v) > 0
), 't_v must be an increasing sequence of time values' # TODO: allow strictly decreasing time values
assert t_v.shape[0] == qp_v.shape[0]
assert t_v.shape[
0] >= 2, 'no discrete derivative can be defined without at least two time values'
assert qp_v.shape[-2] == 2
# T is the number of time values
T = t_v.shape[0]
# N is the dimension of configuration space.
N = qp_v.shape[-1]
retval = np.ndarray(qp_v.shape, dtype=qp_v.dtype)
# First populate retval with discrete derivatives of q and p.
#print('deviation_form; retval[0,...].shape = {0}'.format(retval[0,...].shape))
#print('deviation_form; retval[1:-1,...].shape = {0}'.format(retval[1:-1,...].shape))
#print('deviation_form; (qp_v[2:] - qp_v[:-2]).shape = {0}'.format((qp_v[2:] - qp_v[:-2]).shape))
#print('deviation_form; (t_v[2:] - t_v[:-2]).shape = {0}'.format((t_v[2:] - t_v[:-2]).shape))
retval[0, ...] = (qp_v[1] - qp_v[0]) / (t_v[1] - t_v[0])
retval[1:-1, ...] = qp_v[2:] - qp_v[:-2]
retval[1:-1, ...] /= (t_v[2:] - t_v[:-2]).reshape(
(T - 2, ) + (1, ) * (len(retval.shape) - 1)
) # The reshaping is necessary for correct broadcasting
retval[-1, ...] = (qp_v[-1] - qp_v[-2]) / (t_v[-1] - t_v[-2])
# Then subtract off the partials of the Hamiltonian in order to get the deviation form
retval[...,
0, :] -= vorpy.apply_along_axes(dH_dp, (-1, ),
(qp_v[..., 0, :], qp_v[..., 1, :]),
output_axis_v=(-1, ),
func_output_shape=(N, ))
retval[...,
1, :] += vorpy.apply_along_axes(dH_dq, (-1, ),
(qp_v[..., 0, :], qp_v[..., 1, :]),
output_axis_v=(-1, ),
func_output_shape=(N, ))
return retval
def do_stuff_1():
import sympy as sp
import vorpy
import vorpy.symbolic
import vorpy.symplectic
np.set_printoptions(precision=20)
# Define the Kepler problem and use it to test the integrator
def phase_space_coordinates():
return np.array(sp.var('x,y,p_x,p_y')).reshape(2, 2)
def K(p):
return np.dot(p.flat, p.flat) / 2
def U(q):
return -1 / sp.sqrt(np.dot(q.flat, q.flat))
def H(qp):
"""Total energy -- should be conserved."""
return K(qp[1, ...]) + U(qp[0, ...])
def p_theta(qp):
"""Angular momentum -- should be conserved."""
x, y, p_x, p_y = qp.reshape(-1)
return x * p_y - y * p_x
# Determine the Hamiltonian vector field of H.
qp = phase_space_coordinates()
X_H = vorpy.symplectic.symplectic_gradient_of(H(qp), qp)
print(f'X_H:\n{X_H}')
print('X_H lambdification')
X_H_fast = vorpy.symbolic.lambdified(X_H,
qp,
replacement_d={
'array': 'np.array',
'dtype=object':
'dtype=np.float64'
},
verbose=True)
print('H lambdification')
H_fast = vorpy.symbolic.lambdified(H(qp),
qp,
replacement_d={'sqrt': 'np.sqrt'},
verbose=True)
print('p_theta lambdification')
p_theta_fast = vorpy.symbolic.lambdified(p_theta(qp), qp, verbose=True)
t_initial = 0.0
qp_initial = np.array([[1.0, 0.0], [0.0, 0.5]])
H_initial = H_fast(qp_initial)
p_theta_initial = p_theta_fast(qp_initial)
print(f'H_initial = {H_initial}')
print(f'p_theta_initial = {p_theta_initial}')
#integrator = RungeKutta_4(vector_field=V, parameter_shape=(2,))
integrator = RungeKuttaFehlberg_4_5(
vector_field=(lambda t, qp: X_H_fast(qp)),
parameter_shape=vorpy.tensor.shape(qp_initial))
t = t_initial
y = qp_initial
dt = 0.01
t_max = 3.0
t_v = [t]
y_v = [np.copy(y)]
ltee_v = [0.0]
while t < t_max:
integrator.set_inputs(t, y)
integrator.step(dt)
t, y = integrator.get_outputs()
t_v.append(t)
y_v.append(np.copy(y))
ltee_v.append(np.sqrt(integrator.ltee_squared))
print(f'ltee_v = {ltee_v}')
# Convert the list of np.ndarray to a full np.ndarray.
qp_t = np.array(y_v)
H_v = vorpy.apply_along_axes(H_fast, (1, 2), (qp_t, ))
H_error_v = vorpy.apply_along_axes(
lambda qp: np.abs(H_fast(qp) - H_initial), (1, 2), (qp_t, ))
#print(f'H_v = {H_v}')
#print(f'H_error_v = {H_error_v}')
p_theta_v = vorpy.apply_along_axes(p_theta_fast, (1, 2), (qp_t, ))
p_theta_error_v = vorpy.apply_along_axes(
lambda qp: np.abs(p_theta_fast(qp) - p_theta_initial), (1, 2),
(qp_t, ))
#print(f'p_theta_v = {p_theta_v}')
#print(f'p_theta_error_v = {p_theta_error_v}')
import matplotlib.pyplot as plt
def plot_stuff():
row_count = 1
col_count = 5
size = 5
fig, axis_vv = plt.subplots(row_count,
col_count,
squeeze=False,
figsize=(size * col_count,
size * row_count))
axis = axis_vv[0][0]
axis.set_title('position')
axis.set_aspect('equal')
axis.plot(qp_t[:, 0, 0], qp_t[:, 0, 1], '.')
axis = axis_vv[0][1]
axis.set_title('x and y')
axis.plot(t_v, qp_t[:, 0, 0], '.')
axis.plot(t_v, qp_t[:, 0, 1], '.')
axis = axis_vv[0][2]
axis.set_title('local trunc. err. est.')
axis.semilogy(t_v, ltee_v, '.')
axis = axis_vv[0][3]
axis.set_title('H error')
axis.semilogy(t_v, H_error_v, '.')
axis = axis_vv[0][4]
axis.set_title('p_theta error')
axis.semilogy(t_v, p_theta_error_v, '.')
fig.tight_layout()
filename = 'runge-kutta-kepler.png'
plt.savefig(filename, bbox_inches='tight')
print('wrote to file "{0}"'.format(filename))
# VERY important to do this -- otherwise your memory will slowly fill up!
# Not sure which one is actually sufficient -- apparently none of them are, YAY!
plt.clf()
plt.close(fig)
plt.close('all')
del fig
del axis_vv
plot_stuff()
def do_three_body_problem() -> None:
def phase_space_coordinates() -> np.ndarray:
return np.array(
sp.var('x_0,y_0,x_1,y_1,x_2,y_2,u_0,v_0,u_1,v_1,u_2,v_2')
).reshape(2, 3, 2)
def K(
p: np.ndarray
) -> typing.Any: # Not sure how to annotate a general scalar
return np.dot(p.flat, p.flat) / 2
def U(q: np.ndarray) -> typing.Any:
total = 0
for i in range(3):
for j in range(i):
delta = q[i, :] - q[j, :]
total += -1 / sp.sqrt(np.dot(delta, delta))
return total
def H(qp: np.ndarray) -> typing.Any:
"""Total energy -- should be conserved."""
return K(qp[1, ...]) + U(qp[0, ...])
def p_theta(qp: np.ndarray) -> typing.Any:
"""Angular momentum -- should be conserved."""
total = 0
for i in range(3):
x, y, u, v = qp[:, i, :].reshape(-1)
total += x * v - y * u
return total
def norm_deltas(qp: np.ndarray) -> np.ndarray:
"""Returns the expression defining the distance between bodies i and j for ij = 01, 02, 12 in that order."""
def norm_delta_ij(i: int, j: int) -> typing.Any:
q = qp[0, ...]
q_i = q[i, :]
q_j = q[j, :]
delta = q_i - q_j
return sp.sqrt(np.dot(delta, delta))
return np.array([
norm_delta_ij(0, 1),
norm_delta_ij(0, 2),
norm_delta_ij(1, 2)
])
def body_vertex_angles(qp: np.ndarray) -> np.ndarray:
"""Returns the angle of the vertex of the shape triangle that body i sits at for i = 0, 1, 2."""
def body_vertex_angle_i(i: int) -> typing.Any:
q = qp[0, ...]
# Determine the index of each of the other two bodies.
j = (i + 1) % 3
k = (i + 2) % 3
q_i = q[i, :]
q_j = q[j, :]
q_k = q[k, :]
delta_ij = q_j - q_i
delta_ik = q_k - q_i
norm_squared_delta_ij = np.dot(delta_ij, delta_ij)
norm_squared_delta_ik = np.dot(delta_ik, delta_ik)
cos_angle = np.dot(delta_ij, delta_ik) / sp.sqrt(
norm_squared_delta_ij * norm_squared_delta_ik)
return sp.acos(cos_angle)
return np.array([
body_vertex_angle_i(0),
body_vertex_angle_i(1),
body_vertex_angle_i(2)
])
def shape_area(qp: np.ndarray) -> typing.Any:
"""Area of the parallelopiped spanned by the segments 01 and 02."""
q = qp[0, ...]
delta_01 = q[1, :] - q[0, :]
delta_02 = q[2, :] - q[0, :]
return np.cross(delta_01, delta_02)[()]
# Determine the Hamiltonian vector field of H.
qp = phase_space_coordinates()
X_H = vorpy.symplectic.symplectic_gradient_of(H(qp), qp.reshape(
2, 6)).reshape(2, 3, 2)
print(f'X_H:\n{X_H}')
replacement_d = {
'array': 'np.array',
'dtype=object': 'dtype=np.float64',
'sqrt': 'np.sqrt',
'acos': 'np.arccos',
}
print('X_H lambdification')
X_H_fast = vorpy.symbolic.lambdified(X_H,
qp,
replacement_d=replacement_d,
verbose=True)
print('H lambdification')
H_fast = vorpy.symbolic.lambdified(H(qp),
qp,
replacement_d=replacement_d,
verbose=True)
print('p_theta lambdification')
p_theta_fast = vorpy.symbolic.lambdified(p_theta(qp),
qp,
replacement_d=replacement_d,
verbose=True)
print('shape_area lambdification')
shape_area_fast = vorpy.symbolic.lambdified(
shape_area(qp), qp, replacement_d=replacement_d, verbose=True)
print('norm_deltas lambdification')
norm_deltas_fast = vorpy.symbolic.lambdified(
norm_deltas(qp), qp, replacement_d=replacement_d, verbose=True)
print('body_vertex_angles lambdification')
body_vertex_angles_fast = vorpy.symbolic.lambdified(
body_vertex_angles(qp),
qp,
replacement_d=replacement_d,
verbose=True)
t_initial = 0.0
qp_initial = np.array([
[[-1.0, 0.0], [0.0, 1.0], [1.0, 0.0]],
[[0.0, -0.3], [-0.8, 0.2], [0.1, -0.2]],
])
H_initial = H_fast(qp_initial)
p_theta_initial = p_theta_fast(qp_initial)
print(f'H_initial = {H_initial}')
print(f'p_theta_initial = {p_theta_initial}')
H_cq = ControlledQuantity(
name='H',
reference_quantity=H_initial,
#global_error_band=RealInterval(1.0e-10, 1.0e-6),
global_error_band=RealInterval(1.0e-8, 1.0e-5),
quantity_evaluator=(lambda t, qp: typing.cast(float, H_fast(qp))),
)
p_theta_cq = ControlledQuantity(
name='p_theta',
reference_quantity=p_theta_initial,
#global_error_band=RealInterval(1.0e-10, 1.0e-6),
global_error_band=RealInterval(1.0e-8, 1.0e-5),
quantity_evaluator=(
lambda t, qp: typing.cast(float, p_theta_fast(qp))),
)
controlled_sq_ltee = ControlledSquaredLTEE(
global_error_band=RealInterval(1.0e-12**2, 1.0e-6**2), )
try:
results = integrate_vector_field(
vector_field=(lambda t, qp: X_H_fast(qp)),
t_initial=t_initial,
y_initial=qp_initial,
t_final=10.0,
controlled_quantity_d={
'H abs error': H_cq,
'p_theta abs error': p_theta_cq,
},
controlled_sq_ltee=controlled_sq_ltee,
)
#print(f'results = {results}')
except SalvagedResultsException as e:
print('got SalvagedResultsException')
results = e.args[0]
finally:
H_v = vorpy.apply_along_axes(H_fast, (1, 2, 3), (results.y_t, ))
#H_error_v = vorpy.apply_along_axes(lambda qp:H_cq.error(0.0,qp), (1,2), (results.y_t,))
#print(f'H_v = {H_v}')
#print(f'H_error_v = {H_error_v}')
p_theta_v = vorpy.apply_along_axes(p_theta_fast, (1, 2, 3),
(results.y_t, ))
#p_theta_error_v = vorpy.apply_along_axes(lambda qp:p_theta_cq.error(0.0,qp), (1,2), (results.y_t,))
#print(f'p_theta_v = {p_theta_v}')
#print(f'p_theta_error_v = {p_theta_error_v}')
shape_area_v = vorpy.apply_along_axes(shape_area_fast, (1, 2, 3),
(results.y_t, ))
norm_deltas_t = vorpy.apply_along_axes(norm_deltas_fast, (1, 2, 3),
(results.y_t, ))
print(f'norm_deltas_t = {norm_deltas_t}')
body_vertex_angles_t = vorpy.apply_along_axes(
body_vertex_angles_fast, (1, 2, 3), (results.y_t, ))
print(f'body_vertex_angles_t = {body_vertex_angles_t}')
sign_change_v = shape_area_v[:-1] * shape_area_v[1:]
zero_crossing_index_v = (sign_change_v < 0)
zero_crossing_v = results.t_v[:-1][zero_crossing_index_v]
print(
f'body_vertex_angles_t[:-1][zero_crossing_index_v,:] = {body_vertex_angles_t[:-1][zero_crossing_index_v,:]}'
)
syzygy_v = np.argmax(
body_vertex_angles_t[:-1][zero_crossing_index_v, :], axis=1)
print(f'syzygy_v = {syzygy_v}')
import matplotlib.pyplot as plt
def plot_stuff() -> None:
row_count = 2
col_count = 4
size = 10
fig, axis_vv = plt.subplots(row_count,
col_count,
squeeze=False,
figsize=(size * col_count,
size * row_count))
axis = axis_vv[0][0]
axis.set_title(
f'position (bodies 0,1,2 are R,G,B resp.)\nresults.succeeded = {results.succeeded}\nresults.failure_explanation_o = {results.failure_explanation_o}'
)
axis.set_aspect('equal')
axis.plot(results.y_t[:, 0, 0, 0],
results.y_t[:, 0, 0, 1],
color='red')
axis.plot(results.y_t[:, 0, 1, 0],
results.y_t[:, 0, 1, 1],
color='green')
axis.plot(results.y_t[:, 0, 2, 0],
results.y_t[:, 0, 2, 1],
color='blue')
axis = axis_vv[1][0]
axis.set_title(
'norm delta ij (color is combination of body color)')
axis.semilogy(results.t_v, norm_deltas_t[:, 0], color='yellow')
axis.semilogy(results.t_v,
norm_deltas_t[:, 1],
color='magenta')
axis.semilogy(results.t_v, norm_deltas_t[:, 2], color='cyan')
for zero_crossing in zero_crossing_v:
axis.axvline(zero_crossing, color='black', alpha=0.5)
#axis = axis_vv[1][0]
#axis.set_title('momentum')
#axis.set_aspect('equal')
#axis.plot(results.y_t[:,1,0], results.y_t[:,1,1], '.')
#axis = axis_vv[1][0]
#axis.set_title('blue:(t,x), orange:(t,y)')
#axis.plot(results.t_v, results.y_t[:,0,:], '.')
##axis.plot(results.t_v, results.y_t[:,0,1], '.')
axis = axis_vv[0][1]
axis.set_title('H abs error (global:blue, local:green)')
axis.semilogy(results.t_v,
results.global_error_vd['H abs error'],
'.',
color='blue')
axis.semilogy(results.t_v,
results.local_error_vd['H abs error'],
'.',
color='green')
axis.axhline(
H_cq.global_error_schedule().global_error_band().inf,
color='green')
axis.axhline(
H_cq.global_error_schedule().global_error_band().sup,
color='red')
axis = axis_vv[1][1]
axis.set_title('p_theta abs error (global:blue, local:green)')
axis.semilogy(results.t_v,
results.global_error_vd['p_theta abs error'],
'.',
color='blue')
axis.semilogy(results.t_v,
results.local_error_vd['p_theta abs error'],
'.',
color='green')
axis.axhline(
p_theta_cq.global_error_schedule().global_error_band().inf,
color='green')
axis.axhline(
p_theta_cq.global_error_schedule().global_error_band().sup,
color='red')
axis = axis_vv[0][2]
axis.set_title(ControlledSquaredLTEE.NAME +
' (global:blue, local:green)')
axis.semilogy(
results.t_v,
results.global_error_vd[ControlledSquaredLTEE.NAME],
'.',
color='blue')
axis.semilogy(
results.t_v,
results.local_error_vd[ControlledSquaredLTEE.NAME],
'.',
color='green')
axis.axhline(controlled_sq_ltee.global_error_schedule().
global_error_band().inf,
color='green')
axis.axhline(controlled_sq_ltee.global_error_schedule().
global_error_band().sup,
color='red')
axis = axis_vv[1][2]
axis.set_title('shape_area (vertical line indicates syzygy)')
axis.plot(results.t_v, shape_area_v) #, '.')
axis.axhline(0)
for zero_crossing in zero_crossing_v:
axis.axvline(zero_crossing, color='black', alpha=0.5)
axis = axis_vv[0][3]
axis.set_title('timestep')
axis.semilogy(results.t_v[1:], results.t_step_v, '.')
axis = axis_vv[1][3]
axis.set_title(
'body vertex angles (bodies 0,1,2 are R,G,B resp.)')
axis.plot(results.t_v, body_vertex_angles_t[:, 0], color='red')
axis.plot(results.t_v,
body_vertex_angles_t[:, 1],
color='green')
axis.plot(results.t_v,
body_vertex_angles_t[:, 2],
color='blue')
for zero_crossing in zero_crossing_v:
axis.axvline(zero_crossing, color='black', alpha=0.2)
#axis = axis_vv[1][3]
#axis.set_title('delta timestep (signed log_10 plot)')
#delta_timestep = np.diff(results.t_step_v)
#axis.plot(results.t_v[1:-1], np.sign(delta_timestep)*np.log10(np.abs(delta_timestep)), '.')
fig.tight_layout()
filename = 'threebody.png'
plt.savefig(filename, bbox_inches='tight')
print('wrote to file "{0}"'.format(filename))
# VERY important to do this -- otherwise your memory will slowly fill up!
# Not sure which one is actually sufficient -- apparently none of them are, YAY!
plt.clf()
plt.close(fig)
plt.close('all')
del fig
del axis_vv
plot_stuff()
def do_kepler_problem() -> None:
def phase_space_coordinates() -> np.ndarray:
return np.array(sp.var('x,y,p_x,p_y')).reshape(2, 2)
def K(
p: np.ndarray
) -> typing.Any: # Not sure how to annotate a general scalar
return np.dot(p.flat, p.flat) / 2
def U(q: np.ndarray) -> typing.Any:
return -1 / sp.sqrt(np.dot(q.flat, q.flat))
def H(qp: np.ndarray) -> typing.Any:
"""Total energy -- should be conserved."""
return K(qp[1, ...]) + U(qp[0, ...])
def p_theta(qp: np.ndarray) -> typing.Any:
"""Angular momentum -- should be conserved."""
x, y, p_x, p_y = qp.reshape(-1)
return x * p_y - y * p_x
# Determine the Hamiltonian vector field of H.
qp = phase_space_coordinates()
X_H = vorpy.symplectic.symplectic_gradient_of(H(qp), qp)
print(f'X_H:\n{X_H}')
print('X_H lambdification')
X_H_fast = vorpy.symbolic.lambdified(X_H,
qp,
replacement_d={
'array': 'np.array',
'dtype=object':
'dtype=np.float64'
},
verbose=True)
print('H lambdification')
H_fast = vorpy.symbolic.lambdified(H(qp),
qp,
replacement_d={'sqrt': 'np.sqrt'},
verbose=True)
print('p_theta lambdification')
p_theta_fast = vorpy.symbolic.lambdified(p_theta(qp), qp, verbose=True)
t_initial = 0.0
qp_initial = np.array([[1.0, 0.0], [0.0, 0.01]])
H_initial = H_fast(qp_initial)
p_theta_initial = p_theta_fast(qp_initial)
print(f'H_initial = {H_initial}')
print(f'p_theta_initial = {p_theta_initial}')
H_cq = ControlledQuantity(
name='H',
reference_quantity=H_initial,
global_error_band=RealInterval(1.0e-10, 1.0e-6),
quantity_evaluator=(lambda t, qp: typing.cast(float, H_fast(qp))),
)
p_theta_cq = ControlledQuantity(
name='p_theta',
reference_quantity=p_theta_initial,
global_error_band=RealInterval(1.0e-10, 1.0e-6),
quantity_evaluator=(
lambda t, qp: typing.cast(float, p_theta_fast(qp))),
)
controlled_sq_ltee = ControlledSquaredLTEE(
global_error_band=RealInterval(1.0e-12**2, 1.0e-6**2), )
try:
results = integrate_vector_field(
vector_field=(lambda t, qp: X_H_fast(qp)),
t_initial=t_initial,
y_initial=qp_initial,
t_final=10.0,
controlled_quantity_d={
'H abs error': H_cq,
'p_theta abs error': p_theta_cq,
},
controlled_sq_ltee=controlled_sq_ltee,
)
#print(f'results = {results}')
except SalvagedResultsException as e:
print('got SalvagedResultsException')
results = e.args[0]
finally:
H_v = vorpy.apply_along_axes(H_fast, (1, 2), (results.y_t, ))
#H_error_v = vorpy.apply_along_axes(lambda qp:H_cq.error(0.0,qp), (1,2), (results.y_t,))
#print(f'H_v = {H_v}')
#print(f'H_error_v = {H_error_v}')
p_theta_v = vorpy.apply_along_axes(p_theta_fast, (1, 2),
(results.y_t, ))
#p_theta_error_v = vorpy.apply_along_axes(lambda qp:p_theta_cq.error(0.0,qp), (1,2), (results.y_t,))
#print(f'p_theta_v = {p_theta_v}')
#print(f'p_theta_error_v = {p_theta_error_v}')
import matplotlib.pyplot as plt
def plot_stuff() -> None:
row_count = 2
col_count = 4
size = 5
fig, axis_vv = plt.subplots(row_count,
col_count,
squeeze=False,
figsize=(size * col_count,
size * row_count))
axis = axis_vv[0][0]
axis.set_title('position')
axis.set_aspect('equal')
axis.plot(results.y_t[:, 0, 0], results.y_t[:, 0, 1], '.')
#axis = axis_vv[1][0]
#axis.set_title('momentum')
#axis.set_aspect('equal')
#axis.plot(results.y_t[:,1,0], results.y_t[:,1,1], '.')
axis = axis_vv[1][0]
axis.set_title('blue:(t,x), orange:(t,y)')
axis.plot(results.t_v, results.y_t[:, 0, :], '.')
#axis.plot(results.t_v, results.y_t[:,0,1], '.')
axis = axis_vv[0][1]
axis.set_title('H abs error')
axis.semilogy(results.t_v, results.error_vd['H abs error'],
'.')
axis.axhline(
H_cq.global_error_schedule().global_error_band().inf,
color='green')
axis.axhline(
H_cq.global_error_schedule().global_error_band().sup,
color='red')
axis = axis_vv[1][1]
axis.set_title('p_theta abs error')
axis.semilogy(results.t_v,
results.error_vd['p_theta abs error'], '.')
axis.axhline(
p_theta_cq.global_error_schedule().global_error_band().inf,
color='green')
axis.axhline(
p_theta_cq.global_error_schedule().global_error_band().sup,
color='red')
axis = axis_vv[0][2]
axis.set_title(ControlledSquaredLTEE.NAME)
axis.semilogy(results.t_v,
results.error_vd[ControlledSquaredLTEE.NAME],
'.')
axis.axhline(controlled_sq_ltee.global_error_schedule().
global_error_band().inf,
color='green')
axis.axhline(controlled_sq_ltee.global_error_schedule().
global_error_band().sup,
color='red')
axis = axis_vv[0][3]
axis.set_title('timestep')
axis.semilogy(results.t_v[1:], results.t_step_v, '.')
axis = axis_vv[1][3]
axis.set_title('delta timestep')
delta_timestep = np.diff(results.t_step_v)
axis.plot(
results.t_v[1:-1],
np.sign(delta_timestep) * np.log10(np.abs(delta_timestep)),
'.')
fig.tight_layout()
filename = 'integrator.png'
plt.savefig(filename, bbox_inches='tight')
print('wrote to file "{0}"'.format(filename))
# VERY important to do this -- otherwise your memory will slowly fill up!
# Not sure which one is actually sufficient -- apparently none of them are, YAY!
plt.clf()
plt.close(fig)
plt.close('all')
del fig
del axis_vv
plot_stuff()
def plot_result(self, result_name, axis_v, *, detrend_Hamiltonian=False):
result = self.result_d[result_name]
N = result['N']
can_plot_phase_space = result['can_plot_phase_space']
dt = result['dt']
t_v = result['t_v']
qp_v = result['qp_v']
H_v = result['H_v']
norm_deviation_form_v = result['norm_deviation_form_v']
norm_error_v = result['norm_error_v']
can_plot_phase_space = N == 1
T = qp_v.shape[0]
assert H_v.shape[0] == T
H_v_reshaped = H_v.reshape(T, -1)
if detrend_Hamiltonian:
# Detrend each H_v, so that the plot is just the deviation from the mean and the min and max are more meaningful
H_v_reshaped -= np.mean(H_v_reshaped, axis=0)
norm_deviation_form_v_reshaped = norm_deviation_form_v.reshape(T, -1)
norm_error_v_reshaped = norm_error_v.reshape(T, -1)
min_H = np.min(H_v_reshaped)
max_H = np.max(H_v_reshaped)
range_H = max_H - min_H
max_norm_deviation_form = np.max(norm_deviation_form_v_reshaped[1:-1])
max_norm_error = np.max(norm_error_v_reshaped)
axis = axis_v[0]
axis.set_title('Hamiltonian{0}; range = {1:.2e}\nmethod: {2}'.format(
' (detrended)' if detrend_Hamiltonian else '', range_H,
result_name))
axis.plot(t_v, H_v_reshaped)
axis.axhline(min_H, color='green')
axis.axhline(max_H, color='green')
sqd_v = vorpy.apply_along_axes(lambda x: np.sum(np.square(x)),
(-2, -1), (qp_v - qp_v[0], ))
sqd_v_reshaped = sqd_v.reshape(T, -1)
axis = axis_v[1]
axis.set_title('{0} : sq dist from initial'.format(result_name))
axis.semilogy(t_v, sqd_v_reshaped)
axis = axis_v[2]
axis.set_title(
'{0} : norm of deviation form; max = {1:.2e}\ndt = {2:.2e}'.format(
result_name, max_norm_deviation_form, dt))
axis.semilogy(t_v[1:-1], norm_deviation_form_v_reshaped[1:-1])
axis.axhline(max_norm_deviation_form, color='green')
axis = axis_v[3]
axis.set_title(
'{0} : norm of error (compare to reference); max = {1:.2e}'.format(
result_name, max_norm_error))
axis.semilogy(t_v, norm_error_v_reshaped)
axis.axhline(max_norm_error, color='green')
# We can only directly plot the phase space if the configuration space is 1-dimensional.
if can_plot_phase_space:
# This is to get rid of the last axis, which has size 1, and therefore can be canonically reshaped away
# via the canonical identification between 1-vectors and scalars.
qp_v_reshaped = qp_v.reshape(T, -1, 2)
print('qp_v.shape = {0}'.format(qp_v.shape))
print('qp_v_reshaped.shape = {0}'.format(qp_v_reshaped.shape))
axis = axis_v[-1]
axis.set_title('phase space\nmethod: {0}'.format(result_name))
axis.set_aspect('equal')
axis.plot(qp_v_reshaped[..., 0], qp_v_reshaped[..., 1])