In [ ]:
# Copyright (c) 2015 by Malte Titze (malte.titze@cern.ch)
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
from sympy import *
# import inspect
import math
from __future__ import division # to make sure that 1/2 = 0.5 just in case...
# To handle a rational p/q, use Rational(p,q)
from IPython.display import Math, display
# the symbol s is the 'time' parameter in this worksheet if nothing else specified
s = Symbol('s')
sf = Symbol('sf')
In [ ]:
def get_symbols(expr):
'''
Get all symbols of expr; determine the p's and q's. The 'time' parameter tp is reserved for
symbol s. The canonical variables should be in the form p_1, p_2, ... and q_1, q_2, ... etc.
The output is in the format
[[q_1, q_2, ...], [p_1, p_2, ...], [s], [[s_00(q_1, q_2, ..., p_1, p_2, ..., s),
s_01(q, p, s)], [s_10(q, p, s), s_11(q, p, s), ...], ... ]]
'''
pqs = list(expr.atoms(Symbol))
qs, ps, tp = [], [], []
for ele in pqs:
if str(ele)[0] == 'q':
qs.append(ele)
if str(ele)[0] == 'p':
ps.append(ele)
if str(ele) in ['s', 't']: # perhaps this can be finetuned a bit more here
tp = [ele]
if tp == []:
# the symbol s is the 'time'-parameter in this worksheet. It has to be defined before.
tp = [s]
# now ps = [p_1, p_2, ...], qs = [q_1, q_2, ...] and tp are determined
# the sijs must be treated as functions of ps and qs. In order to declare functions for sympy with
# variable argument lengths we use exec as workaround. Also, the sijs are dependent on s.
sijs = []
i, j = 0, 0
args = qs[:]
for ele in ps:
args.append(ele)
args.append(tp[0])
args = tuple(args)
for q in qs: # if there is no q, then there are no sij's (...)
sijs.append([])
for p in ps:
command = 'sijs[i].append(Function(\'s_\' + str(i) + str(j))' + str(args) + ')'
exec command
j = j + 1
i = i + 1
return [qs, ps, tp, sijs]
def dot_operator(ham, expr):
'''
Implementation of the dot-operator, defined in my report, and
returns the corresponding expression.
'''
# get all symbols of hamiltonian
[qs, ps, tp, sijs] = get_symbols(ham)
sum1 = 0
j = 0
for p in ps:
try:
q = qs[j]
except IndexError:
break
sum1 = sum1 + expr.diff(p)*ham.diff(q)
j = j + 1
sum2 = sum1
i, j = 0, 0
for pi in ps:
for pj in ps:
try:
sum2 = sum2 + expr.diff(pi)*sijs[i][j]*ham.diff(pj)
# may be out of range for sijs since its 1st index i runs over the qs
except IndexError:
break
j = j + 1
i = i + 1
return sum2
def full_operator(ham, expr):
'''
Implementation of the full operator (dot + s-derivative) according to my
report.
'''
[qs, ps, tp, sijs] = get_symbols(ham)
step1 = -dot_operator(ham, expr) + expr.diff(tp[0])
# now we replace the s-derivatives of the sijs with their spacial derivatives
step2 = step1
for i in range(len(ps)):
for j in range(len(qs)):
sum1 = ham.diff(qs[j])
for l in range(len(ps)):
sum1 = sum1 + ham.diff(ps[l])*sijs[j][l]
sum2_interior = sum1
sum1 = -sum1.diff(qs[i])
sum2 = 0
for k in range(len(ps)):
try:
sum2 = sum2 - sum2_interior.diff(ps[k])*sijs[i][k]
# i may be out of range for sijs, since its 1st index has max range len(qs)
except IndexError:
break
sum3 = 0
for l in range(len(qs)):
try:
sum3 = sum3 + ham.diff(qs[l])*sijs[i][j].diff(ps[l])
# l may be out of range for ps[l]
except IndexError:
break
sum4 = 0
for k in range(len(ps)):
for l in range(len(ps)):
sum4 = sum4 + sijs[i][j].diff(ps[k])*sijs[k][l]*ham.diff(ps[l])
dsij = sum1 + sum2 + sum3 + sum4
step2 = step2.subs(sijs[i][j].diff(tp[0]), dsij)
j = j + 1
i = i + 1
return step2
In [ ]:
def build_h_expr(ham, order):
'''
Build the Hamiltonian-expressions for a given order
'''
op = [ham]
for j in range(order - 1):
op.append(full_operator(ham, op[j]))
# op = [ham, full_operator(ham, ham), ... ] has length
# 'order' (for order >=1) at this point.
return op
def qp_cyclic(op, sf):
'''
Implementation of the (implicit) equations for $\bar q, \bar p$
according to my report.
Input:
op: result from build_h_expr
sf: final position wrt. s-coordinate
Output:
[$\bar q$, $\bar p$] the pair of cyclic (final) coordinates. They are given in the form according
to my report, however one has to exchange p with $\bar p$ in all terms involving potentials of elrad manually,
i. e. the output is
$\bar p$ = p + f_1(q, p)*(sf - s) + f_2(q, p)*(sf - s)**2 + ...
and one has to replace the 'p' inside the functions f_k by $\bar p$ (which then makes the equation implicit).
'''
order = len(op)
# get symbols of the underlying hamiltonian
[qs, ps, tp, sijs] = get_symbols(op[0])
elrad = sf - tp[0]
# construct the implicit equation for $\bar p$.
sump = ps[:]
for i in range(len(ps)):
try:
# The number of q's may not be the same as the number of p's.
# We treat the zero'th power of the full_operator seperately:
sump[i] = sump[i] - elrad*op[0].diff(qs[i]).subs(tp[0], sf)
for j in range(1, order):
sump[i] = sump[i] + Rational(1, math.factorial(j + 1))*(-elrad)**(j + 1)*\
op[j].diff(qs[i]).subs(tp[0], sf)
except IndexError:
break
# construct the explicit equation for $\bar q$.
sumq = qs[:]
for i in range(len(qs)):
try:
# The number of p's may not be the same as the number of q's.
# we treat the zero'th power of the full_operator seperately:
sumq[i] = sumq[i] + elrad*op[0].diff(ps[i]).subs(tp[0], sf)
for j in range(1, order):
sumq[i] = sumq[i] - Rational(1, math.factorial(j + 1))*(-elrad)**(j + 1)*\
op[j].diff(ps[i]).subs(tp[0], sf)
except IndexError:
break
# Drop all derivatives of the sij's:
# We do this by parsing the expression to string and replace the word
# 'Derivative' with 0*Derivative and introducing a dummy variable for all sijs:
tempq = str(sumq).replace('Derivative', '0*Derivative')
tempp = str(sump).replace('Derivative', '0*Derivative')
# since the new string is executed, it forgets about the previous variable definitions.
# Therefore we set all sij's to zero.
for cols in sijs:
for ele in cols:
tempq = tempq.replace(str(ele.subs(tp[0], sf)), '0')
tempp = tempp.replace(str(ele.subs(tp[0], sf)), '0')
exec 'q_result = ' + tempq
exec 'p_result = ' + tempp
return [q_result, p_result]
In [ ]:
# Now lets test some 1D examples ...
p = Symbol("p")
q = Symbol("q")
# -------------------
# Hamiltonian
# -------------------
# harmonic oscillator
K = Symbol("K")
ham1 = Rational(1,2)*p**2 + Rational(1,2)*K*q**2
# harmonic oscillator with driving term
lamb = Symbol("lamb")
ham2 = Rational(1,2)*p**2 + Rational(1,2)*K*(1 + Rational(1,6)*cos(lamb*s))**2*q**2
# 1D envelope equation for transversal rms for long beam (the emittance is assumed to be constant)
# (Sacherer 1971)
Q = Symbol("Q") # Q: space charge dependent constant
ham_env = Rational(1,2)*(p**2 + K*q**2 + 1/q**2) + Q*q
# 1D envelope equation for longitudinal rms of a bunched beam inside a conducting pipe (Sacherer 1971)
ham_env2 = Rational(1,2)*(p**2 + K*q**2 + 1/q**2) - Q/q
# normalized paraxial KV (Kapchinskij-Vladimirskij) envelope equation (see Chen & Davidson 1994)
ham_npkv = Rational(1,2)*(p**2 + K*q**2 + 1/q**2) - Q*log(q)
# KV envenlope equation with driving term
ham_npkv2 = Rational(1,2)*(p**2 + K*(1 + Rational(1,5)*cos(lamb*s))**2*q**2 + 1/q**2) - Q*log(q)
# bending magnet (see my previous talk)
ham_bend = 1 - (1 + K*q)*(sqrt(1 - p**2) + K*q*(1 + Rational(1,2))*q)
ham_bend2 = 1 - (1 + q)*(sqrt(1 - p**2) + q*(1 + Rational(1,2))*q)
ham_bend_tmp = sqrt(1 - p**2)*q
In [ ]:
ham = ham_npkv2
print '\nHamiltonian:\n-----------'
display(Math('\mathcal{H} = ' + latex(ham)))
In [ ]:
order = 2
[qb, pb] = qp_cyclic(build_h_expr(ham, order), sf)
In [ ]:
qf = Symbol("qf")
pf = Symbol("pf")
# remember that pb = p + <something where we have to replace p with pf>, see documentation for qp_cyclic.
# therefore we subtract p from pb before we replace p with pf. This subtraction has to be reversed afterwards and
# finally we have to subtract pf, since we bring all variables to one side of the equation.
p_equation = (pb[0] - p).subs(p, pf) + p - pf
In [ ]:
pnew_all = solvers.solve(p_equation, pf)
In [ ]:
pnew = pnew_all[0] # select one of the solutions... carefull in the case of ham_bend_tmp choose the 1 solution!
qnew = qb[0].subs(p, pnew)
In [ ]:
try_simple = False
if try_simple:
qnew, pnew = simplify(qnew), simplify(pnew)
Math('\\begin{align*}' + \
'q^f &= ' + latex( qnew ) + ' ,' + '\\\\' + \
'p^f &= ' + latex( pnew ) + '\\\\' + \
'\\end{align*}')
In [ ]:
# symplecticity check: compute Jacobi-matrix
M = Matrix([\
[qnew.diff(q), qnew.diff(p)],\
[pnew.diff(q), pnew.diff(p)]\
])
print '\nJacobi-matrix:\n-------------\n'
print M
#Math(latex(M))
In [ ]:
# check
jsym = Matrix([[0, 1], [-1, 0]])
matrixcheck = M*jsym*M.transpose() - jsym
try_simple_mat = True
if try_simple_mat:
matrixcheck = simplify(matrixcheck)
print matrixcheck
In [ ]:
check = True
if check:
check_q = 0.39
check_p = 0.7
check_sf = 3
check_s = 2.44
check_K = 1.1
check_lamb = 0.123
for ele in matrixcheck:
print ele.subs(p, check_p).subs(sf, check_sf).subs(s, check_s).subs(lamb, check_lamb).subs(K, check_K)
print '\nsolutions:\n'
print pnew_all
check_expr = p_equation.subs(pf, pnew)
print check_expr
print ''
print N(check_expr.subs(q, check_q).subs(p, check_p).subs(sf, check_sf).subs(s, check_s).subs(K, check_K)\
.subs(lamb, check_lamb))
# This check seems to fail in case of bending magnets
In [ ]: