def test_OdeSolution():
ts = np.array([0, 2, 5], dtype=float)
s1 = ConstantDenseOutput(ts[0], ts[1], np.array([-1]))
s2 = ConstantDenseOutput(ts[1], ts[2], np.array([1]))
sol = OdeSolution(ts, [s1, s2])
assert_equal(sol(-1), [-1])
assert_equal(sol(1), [-1])
assert_equal(sol(2), [-1])
assert_equal(sol(3), [1])
assert_equal(sol(5), [1])
assert_equal(sol(6), [1])
assert_equal(sol([0, 6, -2, 1.5, 4.5, 2.5, 5, 5.5, 2]),
np.array([[-1, 1, -1, -1, 1, 1, 1, 1, -1]]))
ts = np.array([10, 4, -3])
s1 = ConstantDenseOutput(ts[0], ts[1], np.array([-1]))
s2 = ConstantDenseOutput(ts[1], ts[2], np.array([1]))
sol = OdeSolution(ts, [s1, s2])
assert_equal(sol(11), [-1])
assert_equal(sol(10), [-1])
assert_equal(sol(5), [-1])
assert_equal(sol(4), [-1])
assert_equal(sol(0), [1])
assert_equal(sol(-3), [1])
assert_equal(sol(-4), [1])
assert_equal(sol([12, -5, 10, -3, 6, 1, 4]),
np.array([[-1, 1, -1, 1, -1, 1, -1]]))
ts = np.array([1, 1])
s = ConstantDenseOutput(1, 1, np.array([10]))
sol = OdeSolution(ts, [s])
assert_equal(sol(0), [10])
assert_equal(sol(1), [10])
assert_equal(sol(2), [10])
assert_equal(sol([2, 1, 0]), np.array([[10, 10, 10]]))
def compute_field_lines(Bfield,
nperiods=200,
batch_size=8,
magnetic_axis_radius=1,
max_thickness=0.5,
delta=0.01,
steps_per_period=100):
largest = [0.]
def rhs(phi, rz):
nparticles = rz.shape[0] // 2
while phi >= np.pi:
phi -= 2 * np.pi
rz = rz.reshape((nparticles, 2))
rphiz = np.zeros((nparticles, 3))
rphiz[:, 0] = rz[:, 0]
rphiz[:, 1] = phi
rphiz[:, 2] = rz[:, 1]
Brphiz = np.zeros((nparticles, 3))
for i in range(nparticles):
Brphiz[i, :] = Bfield.B(rphiz[i, 0], rphiz[i, 1], rphiz[i, 2])
rhs_rz = np.zeros(rz.shape)
rhs_rz[:, 0] = rphiz[:, 0] * Brphiz[:, 0] / Brphiz[:, 1]
rhs_rz[:, 1] = rphiz[:, 0] * Brphiz[:, 2] / Brphiz[:, 1]
return rhs_rz.flatten()
from scipy.integrate import solve_ivp, RK45, OdeSolution
from math import pi
res = []
nt = int(steps_per_period * nperiods)
tspan = [0, 2 * pi * nperiods]
t_eval = np.linspace(0, tspan[-1], nt + 1)
i = 0
while (i + 1) * batch_size * delta < max_thickness:
y0 = np.zeros((batch_size, 2))
y0[:, 0] = np.linspace(magnetic_axis_radius + i * batch_size * delta,
magnetic_axis_radius +
(i + 1) * batch_size * delta,
batch_size,
endpoint=False)
t = tspan[0]
solver = RK45(rhs,
tspan[0],
y0.flatten(),
tspan[-1],
rtol=1e-9,
atol=1e-09)
ts = [0]
denseoutputs = []
while t < tspan[-1]:
solver.step()
if solver.t < t + 1e-10: # no progress --> abort
break
t = solver.t
ts.append(solver.t)
denseoutputs.append(solver.dense_output())
if t >= tspan[1]:
odesol = OdeSolution(ts, denseoutputs)
res.append(odesol(t_eval))
print(y0[0, 0], "to", y0[-1, 0], "-> success")
else:
print(y0[0, 0], "to", y0[-1, 0], "-> fail")
# break
i += 1
nparticles = len(res) * batch_size
rphiz = np.zeros((nparticles, nt, 3))
xyz = np.zeros((nparticles, nt, 3))
phi_no_mod = t_eval.copy()
for i in range(nt):
while t_eval[i] >= np.pi:
t_eval[i] -= 2 * np.pi
rphiz[:, i, 1] = t_eval[i]
for j in range(len(res)):
for i in range(nt):
rz = res[j][:, i].reshape((batch_size, 2))
rphiz[j * batch_size:(j + 1) * batch_size, i, 0] = rz[:, 0]
rphiz[j * batch_size:(j + 1) * batch_size, i, 2] = rz[:, 1]
for i in range(nt):
for j in range(nparticles):
xyz[j, i, :] = pp.cyl_to_cart(rphiz[j, i, :])
# absB = np.zeros((nparticles, nt))
# tmp = np.zeros((nt, 3))
# for j in range(nparticles):
# tmp[:] = 0
# cpp.biot_savart_B_only(xyz[j, :, :], gammas, dgamma_by_dphis, biotsavart.coil_currents, tmp)
# absB[j, :] = np.linalg.norm(tmp, axis=1)
return rphiz, xyz #, absB, phi_no_mod[:-1]
def compute_field_lines(biotsavart,
nperiods=200,
batch_size=8,
magnetic_axis_radius=1,
max_thickness=0.5,
delta=0.01,
steps_per_period=100):
def cylindrical_to_cartesian(rphiz):
xyz = np.zeros(rphiz.shape)
xyz[:, 0] = rphiz[:, 0] * np.cos(rphiz[:, 1])
xyz[:, 1] = rphiz[:, 0] * np.sin(rphiz[:, 1])
xyz[:, 2] = rphiz[:, 2]
return xyz
gammas = [coil.gamma for coil in biotsavart.coils]
dgamma_by_dphis = [
coil.dgamma_by_dphi[:, 0, :] for coil in biotsavart.coils
]
largest = [0.]
def rhs(phi, rz):
nparticles = rz.shape[0] // 2
while phi >= np.pi:
phi -= 2 * np.pi
rz = rz.reshape((nparticles, 2))
rphiz = np.zeros((nparticles, 3))
rphiz[:, 0] = rz[:, 0]
rphiz[:, 1] = phi
rphiz[:, 2] = rz[:, 1]
xyz = cylindrical_to_cartesian(rphiz)
Bxyz = np.zeros((nparticles, 3))
cpp.biot_savart_B_only(xyz, gammas, dgamma_by_dphis,
biotsavart.coil_currents, Bxyz)
rhs_xyz = np.zeros((nparticles, 3))
rhs_xyz[:, 0] = Bxyz[:, 0]
rhs_xyz[:, 1] = Bxyz[:, 1]
rhs_xyz[:, 2] = Bxyz[:, 2]
# Two different ways of deriving the rhs.
if False:
rhs_phi = (np.cos(phi) * Bxyz[:, 1] -
np.sin(phi) * Bxyz[:, 0]) / rphiz[:, 0]
rhs_rz = np.zeros((nparticles, 2))
rhs_rz[:, 0] = (rhs_xyz[:, 0] * np.cos(phi) +
rhs_xyz[:, 1] * np.sin(phi)) / rhs_phi
rhs_rz[:, 1] = rhs_xyz[:, 2] / rhs_phi
return rhs_rz.flatten()
else:
B_phi = (np.cos(phi) * Bxyz[:, 1] - np.sin(phi) * Bxyz[:, 0])
B_r = np.cos(phi) * Bxyz[:, 0] + np.sin(phi) * Bxyz[:, 1]
B_z = Bxyz[:, 2]
rhs_rz = np.zeros(rz.shape)
rhs_rz[:, 0] = rphiz[:, 0] * B_r / B_phi
rhs_rz[:, 1] = rphiz[:, 0] * B_z / B_phi
return rhs_rz.flatten()
from scipy.integrate import solve_ivp, RK45, OdeSolution
from math import pi
res = []
nt = int(steps_per_period * nperiods)
tspan = [0, 2 * pi * nperiods]
t_eval = np.linspace(0, tspan[-1], nt + 1)
i = 0
while (i + 1) * batch_size * delta < max_thickness:
y0 = np.zeros((batch_size, 2))
y0[:, 0] = np.linspace(magnetic_axis_radius + i * batch_size * delta,
magnetic_axis_radius +
(i + 1) * batch_size * delta,
batch_size,
endpoint=False)
t = tspan[0]
solver = RK45(rhs,
tspan[0],
y0.flatten(),
tspan[-1],
rtol=1e-9,
atol=1e-09)
ts = [0]
denseoutputs = []
while t < tspan[-1]:
solver.step()
if solver.t < t + 1e-10: # no progress --> abort
break
t = solver.t
ts.append(solver.t)
denseoutputs.append(solver.dense_output())
if t >= tspan[1]:
odesol = OdeSolution(ts, denseoutputs)
res.append(odesol(t_eval))
print(y0[0, 0], "to", y0[-1, 0], "-> success")
else:
print(y0[0, 0], "to", y0[-1, 0], "-> fail")
# break
i += 1
nparticles = len(res) * batch_size
rphiz = np.zeros((nparticles, nt, 3))
xyz = np.zeros((nparticles, nt, 3))
phi_no_mod = t_eval.copy()
for i in range(nt):
while t_eval[i] >= np.pi:
t_eval[i] -= 2 * np.pi
rphiz[:, i, 1] = t_eval[i]
for j in range(len(res)):
for i in range(nt):
rz = res[j][:, i].reshape((batch_size, 2))
rphiz[j * batch_size:(j + 1) * batch_size, i, 0] = rz[:, 0]
rphiz[j * batch_size:(j + 1) * batch_size, i, 2] = rz[:, 1]
for i in range(nt):
xyz[:, i, :] = cylindrical_to_cartesian(rphiz[:, i, :])
absB = np.zeros((nparticles, nt))
tmp = np.zeros((nt, 3))
for j in range(nparticles):
tmp[:] = 0
cpp.biot_savart_B_only(xyz[j, :, :], gammas, dgamma_by_dphis,
biotsavart.coil_currents, tmp)
absB[j, :] = np.linalg.norm(tmp, axis=1)
return rphiz, xyz, absB, phi_no_mod[:-1]
def solve_with_event(fun, t_span,y0, method=RK45, t_eval=None,
dense_output=False, events=None,vectorized=False, **options):
# below copied from scipy.integrate # Prepare output
t0, tf = float(t_span[0]), float(t_span[1])
if t_eval is not None:
t_eval = np.asarray(t_eval)
if t_eval.ndim != 1:
raise ValueError("`t_eval` must be 1-dimensional.")
if np.any(t_eval < min(t0, tf)) or np.any(t_eval > max(t0, tf)):
raise ValueError("Values in `t_eval` are not within `t_span`.")
d = np.diff(t_eval)
if tf > t0 and np.any(d <= 0) or tf < t0 and np.any(d >= 0):
raise ValueError("Values in `t_eval` are not properly sorted.")
if tf > t0:
t_eval_i = 0
else:
# Make order of t_eval decreasing to use np.searchsorted.
t_eval = t_eval[::-1]
# This will be an upper bound for slices.
t_eval_i = t_eval.shape[0]
# above copied from scipy.integrate #
# below copied from scipy.integrate # Prepare output
if t_eval is None:
ts = [t0]
ys = [y0]
else:
ts = []
ys = []
interpolants = []
events, is_terminal, event_dir = prepare_events(events)
if events is not None:
g = [event(t0, y0) for event in events]
t_events = [[] for _ in range(len(events))]
# above copied from scipy.integrate #
# add line
t_events_last = [np.NaN for _ in range(len(events))]
# below copied from scipy.integrate # Prepare output
else:
t_events = None
# above copied from scipy.integrate #
funWithEvent = lambda t,y: fun(t,y,t_events_last)
solver = method(funWithEvent, t0, y0, tf, vectorized=vectorized, **options)
# below copied from scipy.integrate # Start integration
status = None
while status is None:
message = solver.step()
if solver.status == 'finished':
status = 0
elif solver.status == 'failed':
status = -1
break
t_old = solver.t_old
t = solver.t
y = solver.y
if dense_output:
sol = solver.dense_output()
interpolants.append(sol)
else:
sol = None
if events is not None:
g_new = [event(t, y) for event in events]
active_events = find_active_events(g, g_new, event_dir)
if active_events.size > 0:
if sol is None:
sol = solver.dense_output()
root_indices, roots, terminate = handle_events(
sol, events, active_events, is_terminal, t_old, t)
for e, te in zip(root_indices, roots):
t_events[e].append(te)
# above copied from scipy.integrate #
# add lines
t_events_last[e]= te
funWithEvent = lambda t,y: fun(t,y,t_events_last)
solver = method(funWithEvent, t, y, tf, vectorized=vectorized, **options)
# below copied from scipy.integrate #
if terminate:
status = 1
t = roots[-1]
y = sol(t)
g = g_new
if t_eval is None:
ts.append(t)
ys.append(y)
else:
# The value in t_eval equal to t will be included.
if solver.direction > 0:
t_eval_i_new = np.searchsorted(t_eval, t, side='right')
t_eval_step = t_eval[t_eval_i:t_eval_i_new]
else:
t_eval_i_new = np.searchsorted(t_eval, t, side='left')
# It has to be done with two slice operations, because
# you can't slice to 0-th element inclusive using backward
# slicing.
t_eval_step = t_eval[t_eval_i_new:t_eval_i][::-1]
if t_eval_step.size > 0:
if sol is None:
sol = solver.dense_output()
ts.append(t_eval_step)
ys.append(sol(t_eval_step))
t_eval_i = t_eval_i_new
message = MESSAGES.get(status, message)
if t_events is not None:
t_events = [np.asarray(te) for te in t_events]
if t_eval is None:
ts = np.array(ts)
ys = np.vstack(ys).T
else:
ts = np.hstack(ts)
ys = np.hstack(ys)
if dense_output:
sol = OdeSolution(ts, interpolants)
else:
sol = None
return OdeResult(t=ts, y=ys, sol=sol, t_events=t_events, nfev=solver.nfev,
njev=solver.njev, nlu=solver.nlu, status=status,
message=message, success=status >= 0)
def solve_ivp_switch(sys, t_span, y0, **kwargs):
kwargs_copy = kwargs.copy()
t_cur = t_span[0]
t_end = t_span[1]
y_cur = y0
t = np.array([])
y = np.array([[] for _ in range(y0.shape[0])])
n_system_events = len(sys.event_functions)
t_sys_events = [[] for _ in range(n_system_events)]
y_sys_events = [[] for _ in range(n_system_events)]
# the event functions of the original system
event_functions = sys.event_functions.copy()
# user event function passed as arguments
user_event_idx = []
try:
user_events = kwargs_copy.pop('events')
if not isinstance(user_events, list):
n_user_events = 1
user_event_idx.append(len(event_functions))
event_functions.append(user_events)
else:
n_user_events = len(user_events)
user_event_idx = []
for event in user_events:
user_event_idx.append(len(event_functions))
event_functions.append(event)
t_events = [[] for _ in range(n_user_events)]
y_events = [[] for _ in range(n_user_events)]
except:
n_user_events = 0
t_events = None
y_events = None
# one last event to stop at the exact time instant
event_functions.append(lambda t, y: t - t_end)
event_functions[-1].terminal = 0
event_functions[-1].direction = 1
n_events = n_system_events + n_user_events + 1
try:
dense_output = kwargs_copy.pop('dense_output')
except:
dense_output = False
if dense_output:
ts = np.array([])
interpolants = []
terminate = False
nfev = 0
njev = 0
nlu = 0
while np.abs(t_cur - t_end) > 1e-10 and not terminate:
sol = solve_ivp(sys, [t_cur, t_end * 1.001],
y_cur,
events=event_functions,
dense_output=True,
**kwargs_copy)
nfev += sol['nfev']
njev += sol['njev']
nlu += sol['nlu']
if not sol['success']:
break
t_next = np.inf
ev_idx = None
for i, t_ev in enumerate(sol['t_events']):
if len(t_ev) > 0 and t_ev[-1] != t_cur and np.abs(t_ev[-1] -
t_next) > 1e-10:
t_next = t_ev[-1]
ev_idx = i
if ev_idx is None:
t_next = sol['t'][-1]
y_next = sol['y'][:, -1]
elif ev_idx in user_event_idx:
y_next = sol['sol'](t_next)
t_events[ev_idx - n_system_events].append(t_next)
y_events[ev_idx - n_system_events].append(sol['sol'](t_next))
if event_functions[ev_idx].terminal:
terminate = True
else:
y_next = sol['sol'](t_next)
if ev_idx < n_system_events:
t_sys_events[ev_idx].append(t_next)
y_sys_events[ev_idx].append(y_next)
S = sys.handle_event(ev_idx, t_next, y_next)
if sys.with_variational:
N = sys.n_dim
phi = S @ np.reshape(y_next[N:], (N, N))
y_next[N:] = phi.flatten()
idx, = np.where(sol['t'] < t_next)
t = np.append(t, sol['t'][idx])
t = np.append(t, t_next)
y = np.append(y, sol['y'][:, idx], axis=1)
y = np.append(y, np.array([y_next]).transpose(), axis=1)
if dense_output:
if len(ts) > 0 and ts[-1] == sol['sol'].ts[0]:
ts = np.concatenate((ts, sol['sol'].ts[1:]))
else:
ts = np.concatenate((ts, sol['sol'].ts))
interpolants += sol['sol'].interpolants
t_cur = t_next
y_cur = y_next
if dense_output:
ode_sol = OdeSolution(ts, interpolants)
else:
ode_sol = None
return OdeResult(t=t, y=y, sol=ode_sol, t_events=t_events, y_events=y_events, \
t_sys_events=t_sys_events, y_sys_events=y_sys_events, \
nfev=nfev, njev=njev, nlu=nlu, status=sol['status'], \
message=sol['message'], success=sol['success'])
