r"""Interface to 1D colnew ode solver.
==========
Overview
==========
Here we provide a simplified interface to the `colnew` boundary value
solver. This can solve differential equations for a set of functions
$Y_i(x)$ of the form
.. math::
Y_i^{(m_i)} = f_i(x, Z)
where
.. math::
Z = [&[Y_1, Y_1', \ldots, Y_1^{(m_1-2)}, Y_1^{(m_1-1)}],\\
&[Y_2, Y_2', \ldots, Y_2^{(m_2-1)}],\\
&[\ldots],\\
&[Y_n, Y_n', \ldots, Y_n^{(m_n-3)}, Y_n^{(m_n-2)}, Y_n^{(m_n-2)}]]
is the set of derivatives $Y^{(m)} = \d{Y}(x)/\d{x}^m$. Thus, there are $n$
differential equations of orders $m_1$, $m_2$, $\ldots$, $m_n`$.
.. note:: The order $m$ may also be zero $m=0$ in which case the equation
represents a constraint to be solved. Likewise, if one of the constraints is
that an integrated quantity has a fixed value (total particle number for
example), then you can introduce $N(x)$ through the first order
equation $N'(x) = n(x)$ and give the boundary conditions $N(0)=0$
and $N(x_{\text{max}}) = N$. These two techniques are often used together to
implement the constraint $N$ with a Lagrange multiplier $\mu$ that is entered
as a parameter with a zeroth-order equation.
A first-order equation with $N(0)=0$ can also be used to simply
integrate $N(x)$ using the adaptive mesh.
The `m_star` = `m_1` + `m_2` + ... `m_n` boundary conditions are specified by
the conditions
.. math::
g_j(y_j, Z) = 0
To simplify the interface, it is assumed that the user will give names
to the `n` variables `Y_n` and define the functions `f_(Y_n)`.
Debugging Tips:
1) Try solving a related problem of the same for, but with a known
exact solution! If you give this as a guess, then you should get
the answer. It is also invaluble to use an exact solution to check
all of the functions to make sure there are no coding mistakes.
2) If you get a `SingularCollocationMatrix: Singular collocation matrix in
COLNEW` error, then try commenting out the derivative functions `d_*` by
adding an underscore `_d_*`. This way, `colnew` will compute these using
finite differences. This is slow, but will catch typo's!
3) Set verbosity to a higher value.
4) Backoff from singular points until things work, then approach them.
5) Add print statements to the various functions to print out
parameters and `x` values. You might find that things are getting
singular.
6) Try setting `self.problem_regularity = 1`. This makes colnew more
careful.
7) Note: The plotting function executes in a try block (to prevent
plotting errors from killing the solution). This can catch
:exc:`KeyboardInterrupt`... You might have to press Ctr-C many times while
plotting...
"""
from __future__ import division
__all__ = ['IDict', 'Solution', 'Transform', 'log_transform', 'ODE', 'ODE_1_2']
import collections
import time
from warnings import warn
import numpy as np
from numpy import pi, sqrt, arctanh, tanh, array
from mmf.objects import StateVars, Container, process_vars
from mmf.objects import Required, Excluded, Computed, ClassVar
import scikits.bvp1lg.jacobian
import scikits.bvp1lg as bvp
_tol = 1e-8
[docs]class IDict(dict):
[docs] def __init__(self, *v, **kw):
dict.__init__(self, *v, **kw)
self.callbacks = []
[docs] def notify_update(self):
for c in self.callbacks:
c(self)
[docs]class Solution(StateVars, bvp.colnew.Solution):
"""Represents the solution to an ODE as returned by colnew. The
default version simply ensures that copies are made of the
abscissa because the colnew solution will mutate these."""
_state_vars = [
('ispace', Required,
"Integer info from colnew about solution."),
('fspace', Required,
"Float info from colnew about solution."),
('ncomp', Computed,
"Number of equations."),
('mstar', Computed,
"Number of degrees of freedom."),
('nmesh', Computed,
"Number of mesh points."),
('ode', Required, "ODE object")]
process_vars()
def __init__(self, *varargin, **kwargs):
bvp.colnew.Solution.__init__(self,
ispace=self.ispace,
fspace=self.fspace)
@property
def x(self):
return self.mesh
def __call__(self, x=None):
"""Evaluation the solution at given points.
:returns:
The solution vector, as::
[u_1(x), u_1'(x), ..., u_1^{m_1 - 1}(x), u_2(x), ...
u_{ncomp}^{m_{ncomp} - 1}]
broadcast to ``x``. Shape of the returned array
is x.shape + (mstar, ).
Also makes sure that x does not get mutated by making a
copy, unlike the default version.
"""
if x is None:
x = self.x
return bvp.colnew.Solution.__call__(self, np.array(x)).T
def plot(self, *varagin, **kwargs):
"""Plot the solution"""
x = self.x
import matplotlib.pylab as plt
plt.plot(x, self(x).T, *varagin, **kwargs)
[docs]class ODE(StateVars):
"""Base class to simplify the interface to colnew."""
_state_vars = [
# Required class variables for defining the problem.
('vars', ClassVar(Required),\
"""This is a list of tuples whos 0th element is a string
representing the variable name (i.e. the name of Y_i) and
who's 1st element is the order of the DEQ for that
variable (i.e. m_i). The 2nd element can be an comment.
This will be used in reporting.
The equation describing the behaviour of this variable
should be specified in the function name_m which should
return the m'th derivative of the variable name."""),
('boundary_points', ClassVar(Required),\
"""List of points where boundary conditions are specified. These
must be in increasing order and must correspond with the index `j`
in the functions."""),
# Optional class variables
('is_linear', ClassVar(False),\
"Set to True if the problem is linear"),
('collocation_points', ClassVar(None),\
"List of explicit collocation points"),
('extra_fixed_points', ClassVar(None),\
"List of explicit fixed points."),
('Solution', Solution,\
"Should inherit from :cls:`Solution`"),
# Instance variables
('left', None, "Left endpoint of abscissa. "\
"Specify if not a boundary point."),
('right', None, "Right endpoint of abscissa. "\
"Specify if not a boundary point."),
('maximum_mesh_size', 1000, "Maximum size for dynamic mesh"),
('problem_regularity', 0, "Set to 1 if the problem is "\
"particularly sensitive"),
('verbosity', 0, ""),
('coarsen_initial_guess_mesh', True,\
"If True, allow mesh to be coarsened."),
('idict', Excluded(IDict()),\
"Dictionary in which to store intermediate results "\
"(I.e. for plotting.)"),
('plot', False, "If True, then plot intermediate results."),
('plot_pause', 0, "Number of seconds to wait after plotting."),
('_debug', False,\
"If True, then some debugging information is printed during "\
"the iteration"),
# Excluded temporary variables.
('m_star', Excluded, "Total number of 'degrees of freedom'.\n"\
"Computed in solve from vars and used in some methods"),
('sol', Excluded, "Temporary storage for current solution."),
('ispace_fspace', Excluded,\
"Temporary storage for the working arrays")
]
process_vars()
def boundary_conditions(self, j, mu, **kwargs):
"""This function must return the `j`'th boundary condition(s),
which should be zero if and only if the boundary condition(s)
at `x = self.boundary_point[j]` is satsified. The functions
must also be smooth. The independent variable is listed in
:attr:`boundary_points`, so `x` is not a parameter. This is used
to build the function g() in colnew. Same rules for arguments
as all previous functions.
Here the boundary conditions are mu(x=0) = self._mu0 and
mu(x=1) = 0.0. Thus, self.boundary_points = [0.0, 1.0] and can
be defined here. Often, however, one passes a parameter to
specify one of the points, in which case, self.boundary_points
would be initialized in self.__init__()
"""
raise NotImplementedError
def check_sol(self, sol):
"""This function should check the solution sol, possibly
correct it and return the resulting solution."""
return sol
def solve(self, sol=None):
"""Return solution.
If sol is provided, then use it to get the initial guess.
"""
self.idict['data'] = collections.deque()
self.ispace_fspace = [None, None]
is_linear = self.is_linear
if self.plot and self.plotter not in self.idict.callbacks:
# Plotting does not work if problem is linear
is_linear = False
self.idict.callbacks.append(self.plotter)
n = len(self.vars)
names = [name_m[0] for name_m in self.vars]
degrees = [name_m[1] for name_m in self.vars]
self.m_star = sum(degrees)
try:
tolerances = self.tolerances
except:
tolerances = np.ones(self.m_star)*_tol
initial_guess = None
if hasattr(self, 'guess'):
initial_guess = self._vectorized_guess
if isinstance(sol, self.Solution):
initial_guess = sol
dgsub = None
if hasattr(self, 'dg') and 'dg' not in names:
dgsub = self._dg
sol = bvp.colnew.solve(
left=self.left,
right=self.right,
boundary_points=self.boundary_points,
degrees=degrees,
fsub=self._vectorized_f,
gsub=self._g,
dfsub=self._vectorized_df,
dgsub=dgsub,
initial_guess=initial_guess,
tolerances=tolerances,
maximum_mesh_size=self.maximum_mesh_size,
vectorized=True,
problem_regularity=self.problem_regularity,
is_linear=is_linear,
collocation_points=self.collocation_points,
verbosity=self.verbosity,
extra_fixed_points=self.extra_fixed_points,
coarsen_initial_guess_mesh=\
self.coarsen_initial_guess_mesh,
#ispace_fspace=self.ispace_fspace,
)
self.sol = self.check_sol(self.Solution(
ode=self, ispace=sol.ispace, fspace=sol.fspace))
return self.sol
def _append_data(self, x, u):
"""Append the current set of data (x, u, f) to self.idict and
perform notifications.
We insert (x, f_dict) to idict['data'] where f_dict is a
dictionary with named fields.
"""
res_dict = dict(x=x)
for i, name_m in enumerate(self.vars):
res_dict[name_m[0]] = u[i, :]
#res_dict['d_' + name] = du[i, :]
res = Container(**res_dict)
self.idict['data'].append(res)
try:
if self.is_linear and 0 < len(self.idict.callbacks):
# If the problem is linear, then the variables are not
# actually tabulated, so warn
warn("For linear problems, function values "
"may not be properly tabulated.")
self.idict.notify_update()
except AttributeError:
pass
def _vectorized_guess(self, x):
us = []
dms = []
for i, xx in enumerate(x):
u, dm = self._guess(float(xx))
us.append(u)
dms.append(dm)
return np.transpose(np.asarray(us)), np.transpose(np.asarray(dms))
def _vectorized_f(self, x, u):
"""Vectorized version of f(x)."""
# du[i, :] is the vector of i'th derivatives at each point in
# x.
du = np.array([self._f(float(xx), u[:, i]) for
(i, xx) in enumerate(x)]).transpose()
return du
def _vectorized_df(self, x, u):
"""Vectorized version of df(x).
Tabulation of plot data is done here because f() may be
evaluated many times when computing the jacobian.
"""
dfs = []
for i, xx in enumerate(x):
dfs.append(self._df(float(xx), u[:, i]))
dfs = np.asarray(dfs)
# Data volatile and needs to be copied, hence the np.array's
self._append_data(np.array(x), np.array(u))
return np.swapaxes(np.swapaxes(dfs, 0, 2), 0, 1)
def _f(self, x, z):
"""Return the derivatives of the equations at the specified
abscissa. This translates between our rep and the expected
bvp format of the data which is an array of all the
derivatives in order.
"""
ans = []
if self._debug:
print("x=%g"%x)
for name_m in self.vars:
name, m = name_m[:2]
f = self.__getattribute__("%s_%i"%(name, m))
vars = self._unpack(z)
der = f(x, **vars)
if self._debug:
l = [" ", "d_%s=%g"%(name, der)]
for n in vars:
l.append("%s=%g"%(n, vars[n]))
print(", ".join(l))
ans.append(der)
return ans
def _df(self, x, z):
ans = []
for name_m in self.vars:
name, m = name_m[:2]
try:
df = self.__getattribute__("d_%s_%i"%(name, m))
except AttributeError:
df = None
if df:
ans.append(self._pack(df(x, **self._unpack(z))))
else:
# Must compute numerical derivative
f = self.__getattribute__("%s_%i"%(name, m))
def F(dz):
vars = self._unpack(z + dz)
return f(x, **vars)
ans.append(bvp.jacobian.jacobian(F,z*0.0))
return ans
def _g(self, z):
ans = []
for j in xrange(self.m_star):
ans.append(self.boundary_conditions(j, **self._unpack(z[:, j])))
return ans
def _dg(self, z):
ans = []
for j in xrange(self.m_star):
ans.append(self._pack(self.d_boundary_conditions(j,
**self._unpack(z[:, j]))))
return ans
def _guess(self, x):
ans = self.guess(x)
z = self._pack(ans)
f = []
for name_m in self.vars:
name, m = name_m[:2]
f.append(ans.get(name, 0.0*x))
return z, f
def _unpack(self, z):
"""Unpack the variables z into a dictionary with specified
names."""
ans = {}
c = 0
for name_m in self.vars:
name, m = name_m[:2]
for n in xrange(m):
ans[name] = z[c]
c += 1
name = "d" + name
return ans
def _pack(self, ders):
"""Return the packed derivatives."""
ans = []
c = 0
for name_m in self.vars:
name, m = name_m[:2]
for n in xrange(m):
ans.append(ders.get(name, 0.0))
c += 1
name = "d" + name
return ans
def plotter(self, idict, **kwargs):
"""Default plotter."""
data = idict['data']
try:
import pylab
pylab.ioff()
pylab.clf()
names = []
for name_m in self.vars:
name, m = name_m[:2]
for j in xrange(m):
names.append(name)
pylab.plot(data[-1].x,
getattr(data[-1],name),
'+-', label=name)
if 1 < len(data):
pylab.plot(data[-2].x,
getattr(data[-2], name), ':')
pylab.xlabel("x")
pylab.ylabel(", ".join(names))
pylab.title("Current mesh size = %i"%len(data[-1].x))
pylab.legend()
pylab.ion()
pylab.draw()
except Exception, e:
warn(str(e))
class _Example_Integrator(ODE):
"""Here is a simple example. This returns the antidirivative of
f(x) on the interval [a, b] with F(0) = 0.
dF/dx = f(x)
F(0) = 0
>>> d = _Example_Integrator(f=np.sin, left=0, right=10.0)
>>> s = d.solve()
>>> x = np.linspace(0, 10, 1000)
>>> max(np.abs(s(x)[0] - (1.0 - np.cos(x)))) < _tol
True
"""
# We have changed the default values of left and right here to
# make them required parameters
_state_vars = [
('f', Required, "Fuction to be integrated"),
('left', Required, "Left endpoint"),
('right', Required, "Right endpoint"),
# The following are the class variables that define the
# problem. These are declared in the ODE base class and
# overloaded here. The don't have to be class variables, but
# that saves each instance from having to make and keep a
# separate copy.
('vars', ClassVar([('F', 1, "Integrated function")])),
('boundary_points', ClassVar([0])),
('is_linear', ClassVar(True)),
]
process_vars()
def F_1(self, x, **kwargs):
"""Defines the equation by returning dF/dx"""
return self.f(x)
def boundary_conditions(self, j, F, **kwargs):
return F
class _Example_Integrator_2(ODE):
"""Here is a simple example. This returns the antidirivative of
f(x) on the interval [a, b] with F(0) = 0.
dF/dx = sqrt(1 - F(x)**2)
F(0) = 0
>>> d = _Example_Integrator(f=np.sin, left=0, right=10.0)
>>> s = d.solve()
>>> x = np.linspace(0, 10, 1000)
>>> max(np.abs(s(x)[0] - (1.0 - np.cos(x)))) < _tol
True
"""
# We have changed the default values of left and right here to
# make them required parameters
_state_vars = [
('f', Required, "Fuction to be integrated"),
('left', Required, "Left endpoint"),
('right', Required, "Right endpoint"),
('vars', ClassVar([('F', 1, "Integrated function")])),
('boundary_points', ClassVar([0])),
('is_linear', ClassVar(False)),
]
process_vars()
def F_1(self, x, F, **kwargs):
return np.sqrt(1.0 - F**2)
def boundary_conditions(self, j, F, **kwargs):
return F
class _Example_Simultaneous(ODE):
"""This is an example of solving two equations or mixed order.
We need to specify three boundary points for this system.
f''(x) = -f(x)
g'(x) = g(x)
"""
_state_vars = [
('left', 0.0),
('right', 1.0),
# The following are the class variables that define the
# problem. These are declared in the ODE base class,
('vars', [('f', 2), ('g', 1)]),
('boundary_points', [0.0, 0.0, 0.0]),
]
process_vars()
def f_2(self, x, f, **kw):
return -f
def g_1(self, x, g, **kw):
return g
def boundary_conditions(self, j, f, df, g, **kwargs):
return [f - 0.0, df - 1.0, g - 1.0][j]
def guess(self, x):
return dict(f=np.sin(x), df=np.cos(x), g=np.exp(x))
class _Test_1a(ODE):
"""This class tests for a memory allocation bug found in previous
versions (005aac368bcc tip) of the bvp package.
Simple differential equation:
u'' = u
We solve this on the interval [a, b] with boundary conditions u(a)
and u(b).
"""
_state_vars = [
('a', Required, "Left endpoint"),
('b', Required, "Right endpoint"),
('ua', Required, "Left value"),
('ub', Required, "Right value"),
('vars', ClassVar([('u', 2)])),
('is_linear', ClassVar(True)),
('boundary_points', Computed),
]
process_vars()
def u_2(self, x, u, **kwargs):
return u
def d_u_2(self, x, u, **kwargs):
return dict(u=1.0)
def boundary_conditions(self, j, u, **kwargs):
return [u - self.ua, u - self.ub][j]
def __init__(self, **kwargs):
self.boundary_points = [self.a, self.b]
@classmethod
def _test(cls):
deq = cls(a=0.0, b=1.0, ua=0.0, ub=1.0)
sol = deq.solve()
X = np.linspace(0.01, 1.0, 100)
U = np.sinh(X)/np.sinh(1)
dU = np.cosh(X)/np.sinh(1)
u = sol(X)[0]
du = sol(X)[1]
abs_err = abs(np.hstack((U-u, dU-du)))
rel_err = abs(abs_err/np.hstack((U, dU)))
print("Abs Err: %g"%(max(abs_err)))
print("Rel Err: %g"%(max(rel_err)))
print("Tol: %g"%(_tol))
class _Test_1b(ODE):
"""Second part of the test for a memory allocation bug found in
previous versions (005aac368bcc tip) of the bvp package.
This solves the equation
u'' = u, u(0) = 0, u(2) = u2
It does this by using the interval [1, 2] with boundary conditions
u(2) = u2 and du(1) is specified by solving the problem on [0, 1]
using _Test_1a. This recursive use demonstrates the bug.
"""
_state_vars = [
('plot', False),
('vars', [('u', 2)]),
('boundary_points', [1.0, 2.0]),
]
process_vars()
def u_2(self, x, u, **kwargs):
return u
def boundary_conditions(self, j, u, du, **kwargs):
if 0 == j:
return du - self._du1_u1(u)
else:
return u - self._u2
def _du1_u1(self, u1):
"""Return du(1) based on u(1) bu solving _Test_1a."""
if self._recursive:
deq = _Test_1a(a=0.0, b=1.0, ua=0.0, ub=u1)
self.sol = deq.solve(self.sol)
du1 = self.sol(1.0)[1]
assert abs(du1 - np.cosh(1.0)*u1/np.sinh(1.0)) < 1e-8
else:
du1 = np.cosh(1.0)*u1/np.sinh(1.0)
print (u1, du1)
return du1
def __init__(self, u2, recursive=False):
self._u2 = u2
self.sol = None
self._recursive = recursive
if __name__ == "__main__":
_Test_1b(1.0, recursive=False).solve()
_Test_1b(1.0, recursive=True).solve()
class _Test_ODE_HartreeFockAtom(ODE):
"""This is an example of how to use the ODE class to find the
energies of spherically symmetric atoms in the Hartee-Fock
approximation.
The differential equation is essentially Poisson's equation for
the chemical potential mu(r)
mu'' + 2*mu'/r = 4*pi*alpha*n(r, mu)
Here n[r, mu(r)] is the total charge density at radius r, including
both the nuclear charge and the electronic charge. However, to
simplify the solution, we treat the nucleus as a point source so
that n[mu] depends only on mu.
The boundary conditions are:
mu(r) = mu0
mu(infty) = 0.0
mu0 must be adjusted after a solution is found so that the total
integrated electronic charge is Z to cancel the nuclear charge.
To deal with the infinite boundary condition, we introduce a
change of variables to x in [0, 1) with a typical scale r0
r = r0*arctanh(x)
dr/dx = r0/(1-x*x)
These will be implemented by the "private" methods
_r_x(x) = r(x), _x_r(r) = x(r), _d_r_x(x) = dr/dx(x) etc.
To simplify the equation, we introduce the variable Q = mu' =
dmu/dr (otherwise the transformation of the equations would be
more complicated.)
We now have two first-order equations:
dmu/dx = Q*dr/dx
dQ/dx = (4*pi*alpha*n[mu] - 2*Q/r)*dr/dx
The variable names are thus Q and mu. These are significant and
will appear throughout the code.
"""
vars = [('mu', 1, "Chemical potential for electrons"),
('Q', 1, "Auxilliary variable Q = dmu/dr")]
plot = True # Define if intermediate results
# should be displayed.
def mu_1(self, x, Q, **kwargs):
"""This function must return the value of the highest order
differential equation for the variable 'mu' as specified in
self.vars. The name of the function must match the name
specified in self.vars.
The first argument must be the independent variable, and the
rest of the arguments must have names corresponding to
self.vars, with a 'd' prepended for derivatives. The equation
for dmu/dx depends only on Q, so we do not need other
arguments.
Note: **kwargs must also be included to "eat" additional arguments.
"""
dmu_x = Q*self._d_r_x(x)
return dmu_x
def d_mu_1(self, x, Q, **kwargs):
"""This is optional, but aids computation. It must return the
derivatives of the function self.mu_1() with respect to the
parameters. Again, the name of the function must correspond
to self.vars and the arguments must satisfy the same
conditions as for self.mu_1().
The results should be returned in a dictionary with the keys
having the names of the variables being differentiated with
respect to. For example, dmu_1(Q)/dQ should be given the key
'Q'. Omitted keys will be assumed to be 0.0 (i.e. no
dependence).
"""
dQ = self._d_r_x(x)
return dict(Q=dQ)
def Q_1(self, x, mu, Q, **kwargs):
"""Here is the highest order differential equation for Q. The
same constraints and format on the arguments and function name
apply as for self.mu_1().
"""
r = self._r_x(x)
dr_x = self._d_r_x(x)
dQ_x = (4.0*pi*self._alpha*self._n(mu) - 2.0*Q/r)*dr_x
return dQ_x
def d_Q_1(self, x, mu, Q, **kwargs):
"""Optional derivatives of the function self.Q() with respect to the
parameters. Same rules as for self.d_mu_1()."""
r = self._r_x(x)
dr_x = self._d_r_x(x)
dQ = -2.0/r*dr_x
dmu = 4.0*pi*self._alpha*self._dn(mu)*dr_x
return dict(Q=dQ, mu=dmu)
boundary_points = [0.0, 1.0]
def boundary_conditions(self, j, mu, **kwargs):
"""This function must return the j'th boundary condition, which should
be zero if and only if the boundary condition at x =
self.boundary_point[j] is satsified. The functions must also
be smooth. The independent variable listed in
self.boundary_points, so 'x' is not a parameter. This is used
to build the function g() in colnew. Same rules for arguments
as all previous functions.
Here the boundary conditions are mu(x=0) = self._mu0 and
mu(x=1) = 0.0. Thus, self.boundary_points = [0.0, 1.0] and can
be defined here. Often, however, one passes a parameter to
specify one of the points, in which case, self.boundary_points
would be initialized in self.__init__()
"""
return [mu - self._mu0, mu - 0.0][j]
def d_boundary_conditions(self, j, mu, **kwargs):
"""This is optional, but aids computation. It must return the
derivatives of the j'th self.boundary_conditions with respect
to the parameters. Same rules for arguments as all previous
functions.
The result should be a dictionaries. The keys correspond to
the names of the parameters (as in the other self.d_*
functions). Omitted keys will again be assumed to be 0.0.
"""
return dict(mu=1.0)
def guess(self, x):
"""Optional function that defines the initial guess.
This should return a dictionary with keys corresponding to the
variable names in self.vars and all of their derivatives.
"""
mu = self._mu0*(self._x_max-x)/self._x_max
dmu = -self._mu0/self._x_max
Q = dmu/self._d_r_x(x)
return dict(mu=mu, dmu=dmu, Q=Q)
###################################################################
# The variables and functions defined below here are "private".
# They are specific to this problem, and have nothing to do with
# the ODE interface.
def _r_x(self, x):
"""Return change of varables r(x)."""
return self._r0*np.arctanh(x)
def _d_r_x(self, x):
"""Return dr/dx(x). This is used to implement the change of
variables."""
return np.divide(self._r0, (1.0-x*x))
def __init__(self, mu0=1.0, x_0=0.0, r0=1.0, x_max=1.0):
"""Initialize local variables etc.
Parameters:
mu0 : Chemical potential at origin (in eV).
r0 : Typical distance scale (in eV^{-1}).
"""
self._alpha = 1./137.035999679 # Fine structure constant.
self._m = 510998.92 # Mass of electron in eV
#self._alpha = 1.0
self._m = 1.0
self._r0 = r0
self._mu0 = mu0
self._x_max = x_max
self.boundary_points[0] = x_0
self.boundary_points[1] = x_max
self.maximum_mesh_size = 10000
def _pF2(self, mu, relativistic=False):
"""Return the Fermi momentum squared pF**2 and its derivative:
"""
m = self._m
if relativistic:
pF2 = mu**2 + 2.0*m*mu # Relativistic
dpF2 = 2.0*mu
else:
pF2 = 2.0*m*mu # Non-relativistic
dpF2 = 2.0*m
return pF2, dpF2
def _n(self, mu):
"""This function computes the density of a free Fermi gas (two
components) with chemical potential mu. We use a relativistic
dispersion relationship here, so actually, the chemical
potential mu specified above is mu - m because of the mass
gap.
n = 2*pF**3/6/pi/pi
"""
pF2, dpF2 = self._pF2(mu)
pF = np.sqrt(pF2)
n = 2.0*pF2*pF/6.0/pi/pi
return np.where(mu>0, n, 0.0)
def _dn(self, mu):
"""Computes the derivative dn/dmu.
dpF/dmu = (dpF2/dmu)/(2.0*pF)
dn/dmu = 3*n/pF*dpF/dmu = pF/2.0/pi/pi*(dpF2/dmu)
"""
pF2, dpF2 = self._pF2(mu)
pF = sqrt(pF2)
dn = pF/2.0/pi/pi*dpF2
return np.where(mu>0, dn, 0.0)
[docs]class ODE_1_2(ODE):
"""Simplified ODE interface for a single second order allowing for
transformations.
Original Equation:
y''(x) = f(x, y, y')
Transformation:
x = s(a)
y = t(b)
Transformed system:
Q'(a) = f(s(a), t(b), Q)/s'(a)
b'(a) = s'(a)/t'(b)*Q(a)
This class requires you to define your
Example:
A''(x) = 2*(A(x) + x*A'(x))
A(-1) = 1
A(2) = exp(3)
Solution: A(x) = exp(x**2-1)
"""
class _Test_ODE_RBoltzmann(ODE):
"""Example problem. Here we use the Relativistic Boltzmann
approximation to model the charge distribution of a nucleus. This
approximation is valid as long as cosh(mu/T) << cosh(m/T), which
is true when n < n0*sinh(m/T) where n0 is give below.
dmu' + 2*dmu/r = 4*pi*alpha*n
dmu = T*n'/sqrt(n0**2+n**2)
n0 = T**3/pi**2*int(x**2/cosh(x**2+(m/T)**2))
The boundary conditions are n(R) = nR and n(R_max) = 0.
"""
_m = 510998.92 # Mass of electron in eV
_alpha = 1./137.035999679
vars = [('n', 1), ('dmu', 1)]
plot = True
def n_1(self, r, n, dmu, **kwargs):
d_n = dmu*sqrt(np.exp(2.0*self._ln_n0)+n**2)/self._T
return d_n
def dmu_1(self, r, n, dmu, **kwargs):
d_dmu = 4.0*pi*self._alpha*n - 2.0*dmu/r
return d_dmu
def boundary_conditions(self, j, n, dmu, **kwargs):
return [n-self._nR, n][j]
def __init__(self, R, nR, T, Rmax=np.inf, m=None):
"""All units in eV"""
import scipy.integrate
if m is not None:
self._m = m
self._nR = nR
self._T = T
self._R = R
self._Rmax = Rmax
self.boundary_points = [self._R, self._Rmax]
a = self._m/self._T
def f(y):
return (np.sqrt(2.0*y+y*y/a)*(1.0+y/a)/
(np.exp(y)+np.exp(-y-2*a)))
self._ln_n0 = np.log(2.0/pi/pi*
scipy.integrate.quad(f, 0, np.inf)[0]*
(self._m*self._T)**(3./2.))-a
n_RB = (np.exp(self._ln_n0+a)-np.exp(self._ln_n0-a))/2.0
if nR >= n_RB:
print(
"""Warning: Initial density too large for Boltzmann approximation!.
nR > n0*sinh(m/T):
%g > %g"""%(nR, n_RB))
class _Test_ODE_RBoltzmann_1(_Test_ODE_RBoltzmann):
"""Example problem. Here we use the Relativistic Boltzmann
approximation to model the charge distribution of a nucleus. This
approximation is valid as long as cosh(mu/T) << cosh(m/T), which
is true when n < n0*sinh(m/T) where n0 is give below.
dmu' + 2*dmu/r = 4*pi*alpha*n
dmu = T*n'/sqrt(n0**2+n**2)
n0 = T**3/pi**2*int(x**2/cosh(x**2+(m/T)**2))
The boundary conditions are n(R) = nR and n(R_max) = 0.
This is the same as RBoltzmann, but uses the transformed variable
u = ln(n/nR), thus u' = n'/n, which keeps values reasonable.
The transformed equations are
dmu' + 2*dmu/r = 4*pi*alpha*nR*exp(u)
dmu = T*u'*nR/sqrt(exp(2*ln_n0-2*u)+nR**2)
"""
vars = [('u', 1), ('dmu', 1)]
plot = True
def _n_u(self, u):
return self._nR*np.exp(u)
def u_1(self, r, u, dmu, **kwargs):
d_u = dmu*sqrt(np.exp(2.0*self._ln_n0-2.0*u)/
self._nR**2+1.0)/self._T
return d_u
def dmu_1(self, r, u, dmu, **kwargs):
d_dmu = 4.0*pi*self._alpha*self._n_u(u) - 2.0*dmu/r
return d_dmu
def boundary_conditions(self, j, u, dmu, **kwargs):
return [u, u-self._umax][j]
def __init__(self, R, nR, T, umax=-np.inf, Rmax=np.inf, m=None):
"""All units in eV"""
import scipy.integrate
if m is not None:
self._m = m
self._nR = nR
self._T = T
self._R = R
self._Rmax = Rmax
self._umax = umax
self.boundary_points = [self._R, self._Rmax]
a = self._m/self._T
def f(y):
return (np.sqrt(2.0*y+y*y/a)*(1.0+y/a)/
(np.exp(y)+np.exp(-y-2*a)))
self._ln_n0 = np.log(2.0/pi/pi*
scipy.integrate.quad(f, 0, np.inf)[0]*
(self._m*self._T)**(3./2.))-a
n_RB = (np.exp(self._ln_n0+a)-np.exp(self._ln_n0-a))/2.0
if nR >= n_RB:
print(
"""Warning: Initial density too large for Boltzmann approximation!.
nR > n0*sinh(m/T):
%g > %g"""%(nR, n_RB))
class _Test_ODE_RBoltzmann_2(_Test_ODE_RBoltzmann_1):
"""Relativistic Boltzmann approximation to model the charge
distribution of a nucleus.
Same as _Test_ODE_RBoltzman_1 but with a different boundary
condition for large r is specified by solving the asymptotic
system which has the solution
mu(r) -> C*T*exp(-2*r*sqrt(pi*alpha*n0/T))/r
n(r) -> C*n0*exp(-2*r*sqrt(pi*alpha*n0/T))/r
This gives us the boundary condition:
n'(Rmax) = -2*sqrt(pi*alpha*n0/T)*n(Rmax) - n(Rmax)/Rmax
or, equivalently
u'(Rmax) = -(2*sqrt(pi*alpha*n0/T) + 1/Rmax)
"""
def boundary_conditions(self, j, u, dmu, **kwargs):
du = self.u(r=self._Rmax, u=u, dmu=dmu)
rhs = -(2.0*sqrt(pi*self._alpha/2)*np.exp(self._ln_n0/2)
+1./self._Rmax)
return [u, du-rhs][j]
def __init__(self, R, nR, T, Rmax=np.inf, m=None):
"""All units in eV"""
_Test_ODE_RBoltzmann_1.__init__(self,
R=R, nR=nR, T=T, Rmax=Rmax, m=m,
umax=None)
del self._umax
print("u(Rmax) must be less than u0 = %g",
self._ln_n0 - np.log(self._nR))
def check_sol(self, sol):
"""Here we check that R_max is large enough."""
Rmax = self._Rmax
umax = sol(Rmax)[0]
u0 = self._ln_n0 - np.log(self._nR)
if umax > u0:
print("""Warning: Rmax not large enough to ensure
asymptotic regime reached.
u(Rmax) > u_inf:
%g > %g"""%(umax, u0))
return sol
class _Test_ODE_RBoltzmann_3(_Test_ODE_RBoltzmann_2):
"""Relativistic Boltzmann approximation to model the charge
distribution of a nucleus.
Same as _Test_ODE_RBoltzman_2 but now with a total charge
constraint Q(Rmax) = Z.
Q(r) = 4*pi*int(r**2*n(r), 0..r)
This is specified by adding an additional differential equation
for the total charge Q enclosed in a sphere of radius r, with an
additional boundary condition at R:
Q'(r) = 4*pi*r**2*n(r)
Q(R) = 4*pi*R**3/3
Q(Rmax) = Z
"""
vars = _Test_ODE_RBoltzmann_2.vars + [('Q', 1)]
def Q_1(self, r, u, **kwargs):
if u > 100.0:
print("u too big (%g)! Clipping..."%u)
n = 0
else:
n = self._n_u(u)
dQ = 4.0*pi*r**2*n
return dQ
def boundary_conditions(self, j, u, dmu, Q, **kwargs):
du = self.u(r=self._Rmax, u=u, dmu=dmu)
rhs = -(2.0*sqrt(pi*self._alpha/2)*np.exp(self._ln_n0/2)
+1./self._Rmax)
Qcore = 3.0*pi*self._R**3/3.0
return [Q-Qcore, du-rhs, Q - self._Z][j]
def __init__(self, R, T, Z, Rmax=np.inf, m=None, nR=1.0):
"""All units in eV"""
_Test_ODE_RBoltzmann_2.__init__(self,
R=R, nR=nR, T=T, Rmax=Rmax, m=m)
self._Z = Z
self.boundary_points = [self._R, self._Rmax, self._Rmax]
print("u(Rmax) must be less than u0 = %g",
self._ln_n0 - np.log(self._nR))
def check_sol(self, sol):
"""Here we check that R_max is large enough."""
Rmax = self._Rmax
umax = sol(Rmax)[0]
u0 = self._ln_n0 - np.log(self._nR)
if umax > u0:
print("""Warning: Rmax not large enough to ensure
asymptotic regime reached.
u(Rmax) > u_inf:
%g > %g"""%(umax, u0))
nR = self._n_u(sol(self._R)[0])
a = self._m/self._T
n_RB = (np.exp(self._ln_n0+a)-np.exp(self._ln_n0-a))/2.0
if nR >= n_RB:
print("""Warning: Core density too large for NR
Boltzmann approximation!. n(R) > n0*sinh(m/T):
%g > %g"""%(nR, n_RB))
return sol
class _Test_ODE_Wall(ODE):
"""Example problem. Here we use the 1D approximation with the
Boltzmann solution to deal with the r=inf limit.
"""
vars = [('zeta', 1)]
alpha = 1./137.035999679
plot = True
def zeta_1(self, x, zeta, **kwargs):
"""
zeta'*zeta = 4*pi*alpha*n(mu)
"""
mu = x
return 4.0*pi*self.alpha*self.n(mu)/zeta
def d_zeta_1(self, x, zeta, **kwargs):
"""Derivative of differential equation."""
mu = x
d_zeta = -4.0*pi*self.alpha*self.n(mu)/zeta/zeta
return dict(zeta=d_zeta)
def boundary_conditions(self, j, zeta, **kwargs):
"""Boundary conditions. This function should return an array
that is zero if the boundary conditions are satisfied and
non-zero (but smooth) otherwise. The boundary conditions
correspond to r in self.boundary_points which must also be
specified."""
return [zeta-self.zeta_B][j]
def d_boundary_conditions(self, j, zeta, **kwargs):
"""Derivative of boundary conditions."""
return [dict(zeta=1.0)][j]
def __init__(self, mu_B, Z=1, R=1):
"""
Parameters
----------
r_inf: Location of right boundary condition where mu=0.0
(in units of 1/eV)
Z: Nuclear charge.
R: Nuclear radius (in fm).
"""
self.eV = 1.0
self.cm = 1.0/1.9732696795479530918e-4/self.eV
self.fm = 1e-13*self.cm
self.m = 510998.92*self.eV
self.Z = Z
self.R = R*self.fm
self.r_inf = r_inf/self.eV
self.boundary_points = [0.0, self.r_inf]
self.extra_fixed_points = [self.R]
def n(self, r, mu):
"""Charge density: non-relativistic T=0 version."""
n_n = np.where(r<self.R, -self.Z/(4.0/3.0*pi*self.R**3), 0)
p_f = np.sign(mu)*sqrt(2.0*self.m*np.abs(mu))
n = n_n + p_f**3/3.0/pi**2
return n
def boundary_conditions(self, j, n, zeta, **kwargs):
"""Boundary conditions. This function should return an array
that is zero if the boundary conditions are satisfied and
non-zero (but smooth) otherwise. The boundary conditions
correspond to r in self.boundary_points which must also be
specified.
zeta(0) = 0
n(1) = 0
"""
return [n-self.n_0, n-self.n_1][j]
#return [zeta, n]
def d_boundary_conditions(self, j, n, dn, **kwargs):
"""Derivative of boundary conditions.
The keys indicate which derivatives are non-zero for each
boundary condition. Omitted keys are assumed to be zero.
"""
return dict(n=1.0)
#return [dict(zeta=1.0), dict(n=1.0)][j]
def __init__(self, r0, r_0, r_1, n_0, n_1, Z=1, R=1, A=1, T=1.0):
"""
Parameters
----------
r0: Scale (in units of 1/eV)
Z: Nuclear charge.
R: Nuclear radius (in fm).
A: 1 for full problem, 0 for 1D approx.
"""
self.A = A
self.eV = 1.0
self.cm = 1.0/1.9732696795479530918e-4/self.eV
self.fm = 1e-13*self.cm
self.m = 510998.92*self.eV
self.Z = Z
self.R = R*self.fm
self.r0 = r0/self.eV
self.r_0 = r_0/self.eV
self.r_1 = r_1/self.eV
self.n_0 = n_0
self.n_1 = n_1
self.T = T
self.boundary_points = [self.z(self.r_0), self.z(self.r_1)]
self.extra_fixed_points = [self.z(self.R)]
def N(self, r):
"""Nuclear charge density."""
return 0.0*r
return np.where(r<self.R, -self.Z/(4.0/3.0*pi*self.R**3), 0)
def guess(self, x):
"""Initial guess. Optional."""
z = x
r = self.r(z)
if self.A == 0.0:
C = (1.0 - 2.0*self.A)/2.0/pi/self.alpha*self.T
r0 = sqrt(C/self.n_0) - self.r_0
n = C/(r+r0)**2
dn = -2.0*C/(r+r0)**3
zeta = dn/n
else:
z_0 = self.z(self.r_0)
z_1 = self.z(self.r_1)
Dz = z_1 - z_0
Dn = self.n_1 - self.n_0
m = Dn/Dz
n = m*(z-z_0) + self.n_0
dn = m + 0.0*x
zeta = dn/n
return dict(n=n, zeta=zeta)
class _Test_ODE_nucleus1(ODE):
"""Example problem. Here we model the charge distribution of
non-relativistic electrons about a nucleus.
There is no hope of solving this without any transformations, but
we can demonstrate the solution by making the scales comparable.
>> p = _Test_ODE_nucleus1(r_inf=2.0)
>> p.R = p.r_inf/2.0
>> s = p.solve()
"""
vars = [('mu', 2)]
alpha = 1./137.035999679
plot = True
def mu(self, x, mu, dmu, **kwargs):
"""Poisson equation:
mu'' + 2*mu'/r = 4*pi*alpha*n(r, mu)
"""
r = x
return 4.0*pi*self.alpha*self.n(r, mu) - 2.0*dmu/r
def d_mu(self, x, mu, dmu, **kwargs):
"""Derivative of differential equation."""
r = x
d_mu = 4.0*pi*self.alpha*self.dn(r, mu)
d_dmu = -2.0/r
return dict(mu=d_mu, dmu=d_dmu)
def boundary_conditions(self, j, mu, dmu, **kwargs):
"""Boundary conditions. This function should return an array
that is zero if the boundary conditions are satisfied and
non-zero (but smooth) otherwise. The boundary conditions
correspond to r in self.boundary_points which must also be
specified."""
return [dmu, mu][j]
def d_boundary_conditions(self, j, mu, dmu, **kwargs):
"""Derivative of boundary conditions."""
return [dict(dmu=1.0), dict(mu=1.0)][j]
def __init__(self, r_inf, Z=1, R=1):
"""
Parameters
----------
r_inf: Location of right boundary condition where mu=0.0
(in units of 1/eV)
Z: Nuclear charge.
R: Nuclear radius (in fm).
"""
self.eV = 1.0
self.cm = 1.0/1.9732696795479530918e-4/self.eV
self.fm = 1e-13*self.cm
self.m = 510998.92*self.eV
self.Z = Z
self.R = R*self.fm
self.r_inf = r_inf/self.eV
self.boundary_points = [0.0, self.r_inf]
self.extra_fixed_points = [self.R]
def n(self, r, mu):
"""Charge density: non-relativistic T=0 version."""
n_n = np.where(r<self.R, -self.Z/(4.0/3.0*pi*self.R**3), 0)
p_f = np.sign(mu)*sqrt(2.0*self.m*np.abs(mu))
n = n_n + p_f**3/3.0/pi**2
return n
class _Test_ODE(ODE):
""" Test system.
A'' = A,
B'' = A*(A'*B - exp(-A))
with boundary conditions
A(-2) = exp(-2)
A(1) = e
B(-2) = exp(-exp(-2))
B(1) = exp(-e)
This has solutions A = exp(x), B=exp(-exp(x))
>>> p = _Test_ODE()
>>> s = p.solve()
>>> x = s.mesh
>>> (abs(s(x)[0]-np.exp(x)) < 1e-13).all()
True
>>> (abs(s(x)[1]-np.exp(x)) < 1e-13).all()
True
>>> (abs(s(x)[2]-np.exp(-np.exp(x))) < 1e-13).all()
True
>>> (abs(s(x)[3]+np.exp(x)*np.exp(-np.exp(x))) < 3e-13).all()
True
"""
vars = [('A', 2),
('B', 2)]
def A_2(self, x, A, **kwargs):
"""Function A''(A)"""
return A
def B_2(self, x, A, dA, B, **kwargs):
"""Function B''(A, A', B)"""
return A*(dA*B - np.exp(-A))
def d_A_2(self, x, A, **kwargs):
"""Derivatives of A''(A). Optional."""
return(dict(A=1.0))
def d_B_2(self, x, A, dA, B, **kwargs):
"""Derivatives of B''(A, A', B), Optional."""
_A = (dA*B - np.exp(-A)) + A*np.exp(-A)
_dA = A*B
_B = A*dA
return(dict(A=_A, B=_B, dA=_dA))
boundary_points = [-2.0, -2.0, 1.0, 1.0]
tolerances = [_tol, _tol, _tol, _tol] # Optional
def boundary_conditions(self, j, A, B, **kwargs):
"""Boundary conditions. This function should return an array
that is zero if the boundary conditions are satisfied and
non-zero (but smooth) otherwise. The boundary conditions
correspond to r in self.boundary_points which must also be
specified."""
return [A-np.exp(-2),
B-np.exp(-np.exp(-2)),
A-np.exp(1),
B-np.exp(-np.exp(1))][j]
def d_boundary_conditions(self, j, A, B, **kwargs):
"""Derivatives of the boundary conditions. Optional."""
return [dict(A=1.0), dict(B=1.0), dict(A=1.0), dict(B=1.0)][j]
def guess(self, x):
"""Initial guess. Optional."""
return dict(A=x, B=x, dA=0.0*x)
def _plot_iter(self, x, A, B, **kwargs):
"""This function is called to plot the current solution. Optional.
One may use the variable self._old_plot_data as a dictionary
to store previous results for comparison.
"""
from pylab import ion, ioff, clf, plot, draw, axis, xlabel, ylabel
ioff()
clf()
plot(x, A, '+-b')
plot(x, B, '+-g')
plot(x, np.exp(x), 'b--')
plot(x, np.exp(-np.exp(x)), 'g--')
if self._old_plot_data.has_key('x'):
plot(self._old_plot_data['x'], self._old_plot_data['A'], 'b:')
plot(self._old_plot_data['x'], self._old_plot_data['B'], 'g:')
self._old_plot_data['x'] = x
self._old_plot_data['A'] = A
self._old_plot_data['B'] = B
xlabel("x")
ylabel("A (blue) and B (green)")
axis([-2, 1, -1, 2])
ion()
draw()
import pdb;pdb.set_trace()