I am seeking to find a finite difference solution to the 1D Nonlinear PDE
u_t = u_xx + u(u_x)^2
Code:
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import math
'''
We explore three different numerical methods for solving the PDE, with solution u(x, t),
u_t = u_xx + u(u_x)^2
for (x, t) in (0, 1) . (0, 1/5)
u(x, 0) = 40 * x^2 * (1 - x) / 3
u(0, t) = u(1, t) = 0
'''
M = 30
dx = 1 / M
r = 0.25
dt = r * dx**2
N = math.floor(0.2 / dt)
x = np.linspace(0, 1, M + 1)
t = np.linspace(0, 0.2, N + 1)
U = np.zeros((M + 1, N + 1)) # Initial array for solution u(x, t)
U[:, 0] = 40 * x**2 * (1 - x) / 3 # Initial condition (: for the whole of that array)
U[0, :] = 0 # Boundary condition at x = 0
U[-1, :] = 0 # Boundary condition at x = 1 (-1 means end of the array)
'''
Explicit Scheme - Simple Forward Difference Scheme
'''
for q in range(0, N - 1):
for p in range(0, M - 1):
b = 1 / (1 - 2 * r)
C = r * U[p, q] * (U[p + 1, q] - U[p, q])**2
U[p, q + 1] = b * (U[p, q] + r * (U[p + 1, q + 1] + U[p - 1, q + 1]) - C)
T, X = np.meshgrid(t, x)
fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_surface(T, X, U)
#fig.colorbar(surf, shrink=0.5, aspect=5) # colour bar for reference
ax.set_xlabel('t')
ax.set_ylabel('x')
ax.set_zlabel('u(x, t)')
plt.tight_layout()
plt.savefig('FDExplSol.png', bbox_inches='tight')
plt.show()
The code I use produces the following error:
overflow encountered in double_scalars
C = r * U[p, q] * (U[p + 1, q] - U[p, q])**2
invalid value encountered in double_scalars
U[p, q + 1] = b * (U[p, q] + r * (U[p + 1, q + 1] + U[p - 1, q + 1]) - C)
invalid value encountered in double_scalars
C = r * U[p, q] * (U[p + 1, q] - U[p, q])**2
Z contains NaN values. This may result in rendering artifacts.
surf = ax.plot_surface(T, X, U)
I've looked up these errors and I assume that the square term generates values too small for the dtype. However when I try changing the dtype to account for a larger range of numbers (np.complex128) I get the same error.
The resulting plot obviously has most of its contents missing. So, my question is, what do I do?
Discretisation expression was incorrect.
Should be
for q in range(0, N - 1):
for p in range(0, M - 1):
U[p, q + 1] = r * (U[p + 1, q] - 2 * U[p, q] + U[p - 1, q]) + r * U[p, q] * (U[p + 1, q] - U[p, q])
Related
My function returns 2 different values which I want to utilise in 2 different graphs using Matplotlib. How can I achieve it?
def option_value_european_put(T, m, r, sigma, mu, E):
cost_value_at_initial_t_put = []
portfolio_payoff_put = []
for e in E:
delta_t = T / m
u = (1 + (sigma * math.sqrt(delta_t)) * (math.sqrt(1 + ((mu ** 2) * delta_t) / math.pow(sigma, 2))))
v = 2 - u
option_stock_price_matrix_put = np.zeros((m + 1, m + 1))
sum = 0
k = m
start = m
for i in range(m + 1):
option_stock_price_matrix_put[i][start] = max(
(e - stock_price_binomial_model(
mu, sigma, T, m,
S
)[i][start], 0)
)
for j in range(m - 1, -1, -1):
for i in range(0, j + 1):
v_plus = option_stock_price_matrix_put[i][j + 1]
v_minus = option_stock_price_matrix_put[i + 1][j + 1]
v_t = ((((v_plus - v_minus) / (u - v)) * (1 + r * delta_t)) + (u * v_minus - v * v_plus) / (u - v)) / (
1 + r * delta_t)
option_stock_price_matrix_put[i][j] = v_t
cost_value_at_initial_t_put.append(option_stock_price_matrix_put[0][0])
for i in range(0, m+1):
sum = sum + option_stock_price_matrix_put[k][i]
portfolio_return_average = math.average(sum)
portfolio_payoff_put.append(portfolio_return_average-option_stock_price_matrix_put[0][0] )
return cost_value_at_initial_t_put, portfolio_payoff_put
I want to use cost_value_at_initial_t_put in 1 Matplotlib plot and the other value in another plot. How can I use it?
Supposing that cost_value_at_initial_t_put and portfolio_payoff_cut are both lists you can create subplots:
import matplotlib.pyplot as plt
fig, (ax_cost, ax_payoff) = plt.subplots(nrows=2)
ax_cost.plot(cost_value_at_initial_t_put)
ax_payoff.plot(portfolio_payoff_cut)
So I am trying to implement some numerical methods into python and I am having some issues where all of my functions output more or less the same thing as the regular euler method. I assume this is because I am messing up in some way when I am implementing the method into code.
My pendulum is defined as this:
def func(y,t):
### Simplified the Function to remove friction since it cancelled out
x,v = y[:3],y[3:6]
grav = np.array([0., 0., -9.8 ])
lambd = (grav.dot(x)+v.dot(v))/x.dot(x)
return np.concatenate([v, grav - lambd*x] )
def dF_matrix(y):
n=y.size
dF=np.zeros((6,6))
xp=np.array([y[1],y[2],y[3]])[np.newaxis]
mass=1.
F1=0.
F2=0.
F3=-mass*9.8
F=np.array([F1,F2,F3])[np.newaxis]
phix=2.*y[0]
phiy=2.*y[4]
phiz=2.*y[5]
G=np.array([phix,phiy,phiz])[np.newaxis]
H=2.*np.eye(3)
lambd=(mass*np.dot(xp,np.dot(H,xp.T))+np.dot(F,G.T))/np.dot(G,G.T)
dF[0,3]=1
dF[1,4]=1
dF[2,5]=1
dF[3,0]=(y[0]*F1+2*lambd)/mass
dF[3,1]=(y[0]*F2)/mass
dF[3,2]=(y[0]*F3)/mass
dF[3,3]=phix*y[1]
dF[3,4]=phix*y[2]
dF[3,5]=phix*y[3]
dF[4,0]=(y[4]*F1)/mass
dF[4,1]=(y[4]*F2+2*lambd)/mass
dF[4,2]=(y[4]*F3)/mass
dF[4,3]=phiy*y[1]
dF[4,4]=phiy*y[2]
dF[4,5]=phiy*y[3]
dF[5,0]=(y[5]*F1)/mass
dF[5,1]=(y[5]*F2)/mass
dF[5,2]=(y[5]*F3+2*lambd)/mass
dF[5,3]=phiz*y[1]
dF[5,4]=phiz*y[2]
dF[5,5]=phiz*y[3]
return dF
The functions that I have made to integrate the ODE function are as follows (with help from others in previous a thread):
from scipy.integrate import odeint
from mpl_toolkits import mplot3d
import matplotlib.pyplot as plt
Forward Euler Method
def forward_Euler(function, y_matrix, time):
y = np.zeros((np.size(time), np.size(y_matrix)))
y[0, :] = y_matrix
for i in range(len(time) - 1):
dt = time[i + 1] - time[i]
y[i + 1, :] = y[i, :] + np.asarray(function(y[i, :], time[i])) * dt
return y
Modified Euler Method
ERROR STARTS HERE
The error I am getting is:
RuntimeWarning: invalid value encountered in double_scalars
lambd = (grav.dot(x)+v.dot(v))/x.dot(x)
def modified_Euler(function, y_matrix, time):
y = np.zeros((np.size(time), np.size(y_matrix))) # creates the matrix that we will fill
y[0, :] = y_matrix # sets the initial values of the matrix
for i in range(len(time) - 1): # apply the Euler
dt = time[i + 1] - time[i] # Step size
k1 = np.asarray(function(y[i, :], time[i])*dt)
k2 = np.asarray(function(y[i] + k1, time[i+1])*dt)
y[i + 1, :] = y[i, :] + .5 * (k1 + k2)
return y
Adams-Bashforth 2nd order
def Adams_Bash_2nd(function, y_matrix, time):
y = np.zeros((np.size(time), np.size(y_matrix)))
y[0, :] = y_matrix
dt = time[1] - time[0]
f_0 = function(y[0], time[0])
y[1] = y[0] + dt * f_0
y[1] = y[0] + 0.5*dt * (f_0+function(y[1], time[1]))
for i in range(len(time) - 1):
dt = time[i + 1] - time[i]
f_1 = function(y[i, :], time[i])
f_2 = function(f_1-1, time[i-1])
y[i + 1] = y[i] + 0.5 * dt * (3 * f_1 - f_2)
return y
Adams Bashforth Moulton Method
def Adams_Moulton(function, y_matrix, time):
y = np.zeros((np.size(time), np.size(y_matrix)))
y[0, :] = y_matrix
### predictor formula
for i in range(len(time) - 1):
dt = time[i + 1] - time[i]
f_1 = function(y[i, :], time[i])
f_2 = function(f_1-1, time[i-1])
y[i + 1, :] = y[i, :] + dt * f_1 + ((dt**2)/2) * f_2
### Corrector formula
for i in range(len(time) - 1):
dt = time[i + 1] - time[i]
k_1 = 9 * (function(y[i, :], time[i+1]))
k_2 = 19 * (function(y[i, :], time[i]))
k_3 = 5 * (function(y[i, :], time[i-1]))
k_4 = (function(y[i, :], time[i-2]))
y[i + 1, :] = y[i] + (dt/24) * (k_1 + k_2 - k_3 + k_4)
return y
RK4 step to use in next function
def RK4_step(f,y,t,dt, N=1):
dt /= N;
for k in range(N):
k1=f(y,t)*dt; k2=f(y+k1/2,t+dt/2)*dt; k3=f(y+k2/2,t+dt/2)*dt; k4=f(y+k3,t+dt)*dt;
y, t = y+(k1+2*(k2+k3)+k4)/6, t+dt
return y
Adams-Bashforth Moulton Method 4th order
def Adams_Moulton_4th(function, y_matrix, time):
y = np.zeros((np.size(time), np.size(y_matrix)))
y[0] = y_matrix
### bootstrap steps with 4th order one-step method
dt = time[4] - time[0]
y[4] = RK4_step(function, y[0], time[0], dt, N=4)
y[5] = RK4_step(function, y[4], time[4], dt, N=4)
y[1] = RK4_step(function, y[5], time[5], dt, N=4)
f_m2 = function(y[0], time[0])
f_m1 = function(y[4], time[4])
f_0 = function(y[5], time[5])
f_1 = function(y[1], time[1])
for i in range(3, len(time) - 1):
### predictor formula 4th order [ 55/24, -59/24, 37/24, -3/8 ]
f_m3, f_m2, f_m1, f_0 = f_m2, f_m1, f_0, f_1
y[i + 1] = y[i] + (dt / 24) * (55 * f_0 - 59 * f_m1 + 37 * f_m2 - 9 * f_m3)
f_1 = function(y[i + 1], time[i + 1])
### Corrector formula 4th order [ 3/8, 19/24, -5/24, 1/24 ]
y[i + 1] = y[i] + (dt / 24) * (9 * f_1 + 19 * f_0 - 5 * f_m1 + f_m2)
f_1 = function(y[i + 1], time[i + 1])
return y
I decided to program the way I am testing the functions into a with a function eliminating a good amount of lines from the previous iteration
# initial condition
y0 = np.array([0.0, 1.0, 0.0, 0.8, 0.0, 1.2])
def test_function(test_function):
print(test_function.__name__ + "...")
nt = 2500
time_start = process_time()
# time points
t = np.linspace(0, 25, nt)
# solve ODE
y1 = test_function(func, y0, t)
time_elapsed = (process_time() - time_start)
print('elapsed time', time_elapsed)
# compute residual:
r = y1[:, 0] ** 2 + y1[:, 1] ** 2 + y1[:, 2] ** 2 - 1
rmax1 = np.max(np.abs(r))
print('error', rmax1)
fig = plt.figure()
ax = plt.axes(projection='3d')
ax.plot3D(y1[:, 0], y1[:, 1], y1[:, 2], 'gray')
plt.show()
test_function(odeint)
test_function(forward_Euler)
test_function(modified_Euler)
test_function(Adams_Bash_2nd)
test_function(Adams_Moulton)
test_function(Adams_Moulton_4th)
The modified Euler method Does Not access points outside the step i -> i+1, there is no i-1 (note that in your source document the step, in the python code, not the formulas, is i-1 -> i with the loops starting at an appropriately increased index). It simply is (as you can find everywhere the mod. Euler or Heun method is discussed)
k1 = f(y[i] , t[i ])*dt;
k2 = f(y[i]+k1, t[i+1])*dt;
y[i+1] = y[i] + 0.5*(k1+k2);
In contrast, the Adams-Bashford method of order 2 and Adams-Moulton methods of order greater 2 Do access points from before the step i -> i+1, formally one has in AB2
y[i+1] = y[i] + 0.5*dt * (3*f[i] - f[i-1])
For a first implementation it would make sense to declare the f array the same way as the y array to implement this formula verbatim. It can be more economical to only keep a short array of f values that is shifted or rotated to give access to the last few f values.
Note that you need to initialize y[1] and f[1] with some other method of similar or higher order. Or if you want to have a "pure" run of the method, you need to initialize y[-1] and f[-1] and further back so that y[1] can be computed with the method formula.
I am intending to take a list of random variables and alter a previous list in each column by said random variables. However, for the purpose of my function, each variable must be used in a Gamma function as well as integrated.
x[t] = c * (1 / (2 ** (v / 2) + test[t - 1]) * (gamma((v / 2) + test[t - 1]))) * integrate.\
quad(lambda h: np.exp(-h / 2) * h ** ((v / 2) + test[t - 1] - 1), 0, np.inf)
x[ t ] is an np.zeros((x , y)) list, and test[t - 1] is an np.zeros((x - 1, y)) list
I have filled test[ ] with the appropriate random variables, but I am unable to pass them through this equation to complete the columns of row [ t ] in x
When I try to run my current code, I receive:
File "C:\Program Files (x86)\Microsoft Visual Studio\Shared\Python37_64\lib\site-packages\scipy\integrate\quadpack.py", line 450, in _quad
return _quadpack._qagie(func,bound,infbounds,args,full_output,epsabs,epsrel,limit)
TypeError: only size-1 arrays can be converted to Python scalars
Is there a different special function which allows me to use each column's variable to solve for my desired x[ t ]?
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import stats
import mpmath as mp
import scipy.integrate as integrate
from scipy.special import gamma
T = 1
beta = 0.5
x0 = 0.05
q = 0
mu = x0 - q
alpha = - (2 - beta) * mu
sigma0 = 0.1
sigma = (2 - beta) * sigma0
b = - ((1 - beta) / (2 * mu) * (sigma0 ** 2))
simulations = 100
M = 50
dt = T / M
def srd_sampled_nxc2():
x = np.zeros((M + 1, simulations))
x[0] = x0
test = np.zeros((M, simulations))
for t in range(1, M + 1):
v = 4 * b * alpha / sigma ** 2
c = (sigma ** 2 * (1 - np.exp(-alpha * dt))) / (4 * alpha)
nc = np.exp(-alpha * dt) / c * x[t - 1]
if v > 1:
x[t] = c * ((np.random.standard_normal(simulations) + nc ** 0.5) ** 2 + mp.nsum(
lambda i: np.random.standard_normal(simulations) ** 2, [0, v - 1]))
else:
max_array = []
nc_over_2 = [l / 2 for l in nc]
for p in range(simulations):
sump = []
poisson_start = 0
while poisson_start <= 1:
x_i = sum(-np.log(np.random.uniform(0, 1, simulations)) / nc_over_2)
sump.append(
x_i
)
poisson_start += x_i
x_n = max(sump)
max_array.append(
x_n
)
sump = []
test[t - 1] = max_array
x[t] = c * (1 / (2 ** ((v / 2) + test[t - 1])) * (gamma((v / 2) + test[t - 1]))) * integrate.\
quad(lambda h: np.exp(-h / 2) * h ** ((v / 2) + test[t - 1] - 1), 0, np.inf)
max_array = []
return x
Ultimately ended up finding a workaround which is simple to implement:
else:
max_array = []
for p in range(simulations):
k = nc[t - 1, p]
lam = k / 2
poisson_samp = 0
while poisson_samp <= 1:
x_i = -math.log(np.random.uniform(0, 1)) / lam
max_array.append(
x_i
)
poisson_samp += x_i
test[t - 1, p] = len(max_array) - 1
max_array.clear()
for f in range(simulations):
n = test[t - 1, f]
z = integrate.quad(lambda h: np.exp(-h / 2) * h ** ((v / 2) + n - 1), 0, 1)
new[t - 1, f] = z[0]
x[t] = c * (1 / (2 ** ((v / 2) + test[t - 1]) * (gamma((v / 2) + test[t - 1]))) * new[0])
The only real problem is the shrinkage of x[t] which leads to dividing by zero--just a formula problem.
I am new to python and numpy there is this Second order PDE code which i want to vectorize to run in lesser time but since it uses function to fill each value in grid i m stuck.
def func(x, y):
return (-2 * np.pi ** 2) * np.sin(np.pi * x) * np.sin(np.pi * y)
def f1t(x, y):
return (np.sin(np.pi * x) * np.sin(np.pi * y))
def five_point(grid, i, j, h, grid_x, grid_y):
return ((grid[i + 1, j] + grid[i - 1, j] + grid[i, j + 1] + grid[i, j - 1]) / 4
- ((h ** 2) / 4) * func(grid_x[i, 0], grid_y[0, j]))
def five_point_fin_int(X=np.ones(1), Y=np.ones(1), n_x=32, K_max=1000,
tol=0.0001, tol_type="grid"):
import time;
t0 = time.clock()
h = 1 / n_x
X_max = int(X / h)
Y_max = int(Y / h)
grid_x, grid_y = np.mgrid[0:X + h:h, 0:Y + h:h]
grid_true = np.zeros((X_max + 1, Y_max + 1))
for i in range(1, X_max):
for j in range(1, Y_max):
grid_true[i, j] = f1t(grid_x[i, 0], grid_y[0, j])
grid = np.zeros((X_max + 1, Y_max + 1))
grid_err = np.zeros((X_max + 1, Y_max + 1))
count = 0
tol_max = False
while ((count < K_max) & (tol_max == False)):
count += 1
for i in range(1, X_max):
for j in range(1, Y_max):
grid[i, j] = five_point(grid, i, j, h, grid_x, grid_y)
grid_err[i, j] = (grid[i, j] - grid_true[i, j])
if (tol_type.lower() == "grid" ):
if (np.amax(abs(grid_err)) < tol):
tol_max = True
else:
if (abs(np.linalg.norm(grid) - np.linalg.norm(grid_true)) < tol):
tol_max = True
cpu_time = time.clock() - t0
In the end i print compute time since its nested for loops right now the time taken is a lot around 9 sec i would like to improvise on this.
numpy allows you to replace loops by vector calls. You can definitely do the following:
grid_true = np.zeros((X_max + 1, Y_max + 1))
grid_true[1:X_max,1:Y_max]=f1t(*np.meshgrid(grid_x[1:X_max,0], grid_y[0,1:Y_max]))
And you can also try the following:
grid = np.zeros((X_max + 1, Y_max + 1))
grid[1:-1, 1:-1] = five_point(grid, *np.meshgrid(np.arange(1,X_max), np.arange(1, Y_max)),
h, grid_x, grid_y)
However this is not pure "upstream" integration like the one you are doing, since you are essentialy calculating all the grid together in each step (your call!).
Probably a minimization routine could do better. There isn't much difference in performance between numpy and pure python for short loops or small vectors.
For a fixed integer n, I have a set of 2(n-1) simultaneous equations as follows.
M(p) = 1+((n-p-1)/n)*M(n-1) + (2/n)*N(p-1) + ((p-1)/n)*M(p-1)
N(p) = 1+((n-p-1)/n)*M(n-1) + (p/n)*N(p-1)
M(1) = 1+((n-2)/n)*M(n-1) + (2/n)*N(0)
N(0) = 1+((n-1)/n)*M(n-1)
M(p) is defined for 1 <= p <= n-1. N(p) is defined for 0 <= p <= n-2. Notice also that p is just a constant integer in every equation so the whole system is linear.
I have been using Maple but I would like to set these up and solve them in python now, maybe using numpy.linalg.solve (or any other better method). I actually only want the value of M(n-1). For example, when n=2 the answer should be M(1) = 4, I believe. This is because the equations become
M(1) = 1+(2/2)*N(0)
N(0) = 1 + (1/2)*M(1)
Therefore
M(1)/2 = 1+1
and so
M(1) = 4.
If I want to plug in n=50, say, how can you set up this system of simultaneous equations in python so that numpy.linalg.solve can solve them?
Update The answers are great but they use dense solvers where the system of equations is sparse. Posted follow up to Using scipy sparse matrices to solve system of equations .
Updated: added implementation using scipy.sparse
This gives the solution in the order N_max,...,N_0,M_max,...,M_1.
The linear system to solve is of the shape A dot x == const 1-vector.
x is the sought after solution vector.
Here I ordered the equations such that x is N_max,...,N_0,M_max,...,M_1.
Then I build up the A-coefficient matrix from 4 block matrices.
Here's a snapshot for the example case n=50 showing how you can derive the coefficient matrix and understand the block structure. The coefficient matrix A is light blue, the constant right side is orange. The sought after solution vector x is here light green and used to label the columns. The first column show from which of the above given eqs. the row (= eq.) has been derived:
As Jaime suggested, multiplying by n improves the code. This is not reflected in the spreadsheet above but has been implemented in the code below:
Implementation using numpy:
import numpy as np
import numpy.linalg as linalg
def solve(n):
# upper left block
n_to_M = -2. * np.eye(n-1)
# lower left block
n_to_N = (n * np.eye(n-1)) - np.diag(np.arange(n-2, 0, -1), 1)
# upper right block
m_to_M = n_to_N.copy()
m_to_M[1:, 0] = -np.arange(1, n-1)
# lower right block
m_to_N = np.zeros((n-1, n-1))
m_to_N[:,0] = -np.arange(1,n)
# build A, combine all blocks
coeff_mat = np.hstack(
(np.vstack((n_to_M, n_to_N)),
np.vstack((m_to_M, m_to_N))))
# const vector, right side of eq.
const = n * np.ones((2 * (n-1),1))
return linalg.solve(coeff_mat, const)
Solution using scipy.sparse:
from scipy.sparse import spdiags, lil_matrix, vstack, hstack
from scipy.sparse.linalg import spsolve
import numpy as np
def solve(n):
nrange = np.arange(n)
diag = np.ones(n-1)
# upper left block
n_to_M = spdiags(-2. * diag, 0, n-1, n-1)
# lower left block
n_to_N = spdiags([n * diag, -nrange[-1:0:-1]], [0, 1], n-1, n-1)
# upper right block
m_to_M = lil_matrix(n_to_N)
m_to_M[1:, 0] = -nrange[1:-1].reshape((n-2, 1))
# lower right block
m_to_N = lil_matrix((n-1, n-1))
m_to_N[:, 0] = -nrange[1:].reshape((n-1, 1))
# build A, combine all blocks
coeff_mat = hstack(
(vstack((n_to_M, n_to_N)),
vstack((m_to_M, m_to_N))))
# const vector, right side of eq.
const = n * np.ones((2 * (n-1),1))
return spsolve(coeff_mat.tocsr(), const).reshape((-1,1))
Example for n=4:
[[ 7.25 ]
[ 7.76315789]
[ 8.10526316]
[ 9.47368421] # <<< your result
[ 9.69736842]
[ 9.78947368]]
Example for n=10:
[[ 24.778976 ]
[ 25.85117842]
[ 26.65015984]
[ 27.26010007]
[ 27.73593401]
[ 28.11441922]
[ 28.42073207]
[ 28.67249606]
[ 28.88229939]
[ 30.98033266] # <<< your result
[ 31.28067182]
[ 31.44628982]
[ 31.53365219]
[ 31.57506477]
[ 31.58936225]
[ 31.58770694]
[ 31.57680467]
[ 31.560726 ]]
Here's an entirely different approach, using sympy. It's not fast, but it allows me to copy the RHS of your equations exactly, limiting the thinking I need to do (always a plus), and gives fractional answers.
from sympy import Integer, Symbol, Eq, solve
def build_equations(n):
ni = n
n = Integer(n)
Ms = {p: Symbol("M{}".format(p)) for p in range(ni)}
Ns = {p: Symbol("N{}".format(p)) for p in range(ni-1)}
M = lambda i: Ms[int(i)] if i >= 1 else 0
N = lambda i: Ns[int(i)]
M_eqs = {}
M_eqs[1] = Eq(M(1), 1+((n-2)/n)*M(n-1) + (2/n)*N(0))
for p in range(2, ni):
M_eqs[p] = Eq(M(p), 1+((n-p-1)/n)*M(n-1) + (2/n)*N(p-1) + ((p-1)/n)*M(p-1))
N_eqs = {}
N_eqs[0] = Eq(N(0), 1+((n-1)/n)*M(n-1))
for p in range(1, ni-1):
N_eqs[p] = Eq(N(p), 1+((n-p-1)/n)*M(n-1) + (p/n)*N(p-1))
return M_eqs.values() + N_eqs.values()
def solve_system(n, show=False):
eqs = build_equations(n)
sol = solve(eqs)
if show:
print 'equations:'
for eq in sorted(eqs):
print eq
print 'solution:'
for var, val in sorted(sol.items()):
print var, val, float(val)
return sol
which gives
>>> solve_system(2, True)
equations:
M1 == N0 + 1
N0 == M1/2 + 1
solution:
M1 4 4.0
N0 3 3.0
{M1: 4, N0: 3}
>>> solve_system(3, True)
equations:
M1 == M2/3 + 2*N0/3 + 1
M2 == M1/3 + 2*N1/3 + 1
N0 == 2*M2/3 + 1
N1 == M2/3 + N0/3 + 1
solution:
M1 34/5 6.8
M2 33/5 6.6
N0 27/5 5.4
N1 5 5.0
{M2: 33/5, M1: 34/5, N1: 5, N0: 27/5}
and
>>> solve_system(4, True)
equations:
M1 == M3/2 + N0/2 + 1
M2 == M1/4 + M3/4 + N1/2 + 1
M3 == M2/2 + N2/2 + 1
N0 == 3*M3/4 + 1
N1 == M3/2 + N0/4 + 1
N2 == M3/4 + N1/2 + 1
solution:
M1 186/19 9.78947368421
M2 737/76 9.69736842105
M3 180/19 9.47368421053
N0 154/19 8.10526315789
N1 295/38 7.76315789474
N2 29/4 7.25
{N2: 29/4, N1: 295/38, M1: 186/19, M3: 180/19, N0: 154/19, M2: 737/76}
which seems to match the other answers.
This is messy, but solves your problem, barring a very probable mistake transcribing the coefficients:
from __future__ import division
import numpy as np
n = 2
# Solution vector is [N[0], N[1], ..., N[n - 2], M[1], M[2], ..., M[n - 1]]
n_pos = lambda p : p
m_pos = lambda p : p + n - 2
A = np.zeros((2 * (n - 1), 2 * (n - 1)))
# p = 0
# N[0] + (1 - n) / n * M[n-1] = 1
A[n_pos(0), n_pos(0)] = 1 # N[0]
A[n_pos(0), m_pos(n - 1)] = (1 - n) / n #M[n - 1]
for p in xrange(1, n - 1) :
# M[p] + (1 + p - n) /n * M[n - 1] - 2 / n * N[p - 1] +
# (1 - p) / n * M[p - 1] = 1
A[m_pos(p), m_pos(p)] = 1 # M[p]
A[m_pos(p), m_pos(n - 1)] = (1 + p - n) / n # M[n - 1]
A[m_pos(p), n_pos(p - 1)] = -2 / n # N[p - 1]
if p > 1 :
A[m_pos(p), m_pos(p - 1)] = (1 - p) / n # M[p - 1]
# N[p] + (1 + p -n) / n * M[n - 1] - p / n * N[p - 1] = 1
A[n_pos(p), n_pos(p)] = 1 # N[p]
A[n_pos(p), m_pos(n - 1)] = (1 + p - n) / n # M[n - 1]
A[n_pos(p), n_pos(p - 1)] = -p / n # N[p - 1]
if n > 2 :
# p = n - 1
# M[n - 1] - 2 / n * N[n - 2] + (2 - n) / n * M[n - 2] = 1
A[m_pos(n - 1), m_pos(n - 1)] = 1 # M[n - 1]
A[m_pos(n - 1), n_pos(n - 2)] = -2 / n # N[n - 2]
A[m_pos(n - 1), m_pos(n - 2)] = (2 - n) / n # M[n - 2]
else :
# p = 1
#M[1] - 2 / n * N[0] = 1
A[m_pos(n - 1), m_pos(n - 1)] = 1
A[m_pos(n - 1), n_pos(n - 2)] = -2 / n
X = np.linalg.solve(A, np.ones((2 * (n - 1),)))
But it gives a solution of
>>> X[-1]
6.5999999999999979
for M(2) when n=3, which is not what you came up with.