From 34fdd97b12e1a3686e79871c1527d80076bb57bb Mon Sep 17 00:00:00 2001 From: Mark Dewing Date: Tue, 29 Jan 2019 18:51:24 -0600 Subject: [PATCH] Basic cubic spline coefficient solver --- Wavefunctions/Cubic Splines Basic.ipynb | 1263 +++++++++++++++++++++++ Wavefunctions/eqn_manip.py | 49 + 2 files changed, 1312 insertions(+) create mode 100644 Wavefunctions/Cubic Splines Basic.ipynb create mode 100644 Wavefunctions/eqn_manip.py diff --git a/Wavefunctions/Cubic Splines Basic.ipynb b/Wavefunctions/Cubic Splines Basic.ipynb new file mode 100644 index 0000000..01a7571 --- /dev/null +++ b/Wavefunctions/Cubic Splines Basic.ipynb @@ -0,0 +1,1263 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "from __future__ import print_function\n", + "from sympy import *\n", + "from eqn_manip import move_terms, move_terms_left, divide_terms, mult_eqn, get_coeff_for\n", + "from IPython.display import display\n", + "init_printing()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Cubic spline equations\n", + "First step is make the a straightforward derivation of some cubic spline equations. These could be further simplified to a tridiagonal matrix and efficiently solved. For small numbers of knots and for testing purposes, it is sufficient to use a larger matrix and a general solver.\n", + "\n", + "The first derivation will involve knots with constant separation of 1. Later the derivation will be repeated with arbitrary knot positions." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$$y = t^{3} d + t^{2} c + t b + a$$" + ], + "text/plain": [ + " 3 2 \n", + "y = t ⋅d + t ⋅c + t⋅b + a" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/latex": [ + "$${y}_{i} = t^{3} {d}_{i} + t^{2} {c}_{i} + t {b}_{i} + {a}_{i}$$" + ], + "text/plain": [ + " 3 2 \n", + "y[i] = t ⋅d[i] + t ⋅c[i] + t⋅b[i] + a[i]" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "# Distance from the previous knot, for the case of uniform knot spacing\n", + "t = Symbol('t')\n", + "\n", + "# Number of knots\n", + "n = Symbol('n', integer=True)\n", + "i = Symbol('i', integer=True)\n", + "# Function values to intepolated at the knots\n", + "y = IndexedBase('y')\n", + "\n", + "# Coefficients of the spline function\n", + "a,b,c,d = [IndexedBase(s) for s in 'a b c d'.split()]\n", + "\n", + "# Cubic spline function\n", + "s = a + b*t + c*t*t + d*t**3\n", + "display(Eq(y,s))\n", + "\n", + "# With indexed variables\n", + "si = a[i] + b[i]*t + c[i]*t*t + d[i]*t**3\n", + "display(Eq(y[i],si))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Set up equations" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$${a}_{i} = {y}_{i}$$" + ], + "text/plain": [ + "a[i] = y[i]" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Value at knots (t=0)\n", + "eq1 = Eq(si.subs(t,0), y[i])\n", + "eq1" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$${a}_{i} + {b}_{i} + {c}_{i} + {d}_{i} = {y}_{i + 1}$$" + ], + "text/plain": [ + "a[i] + b[i] + c[i] + d[i] = y[i + 1]" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Value at knots (t=1)\n", + "eq2 = Eq(si.subs(t,1), y[i+1])\n", + "eq2" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$${b}_{i} + 2 {c}_{i} + 3 {d}_{i} = {b}_{i + 1}$$" + ], + "text/plain": [ + "b[i] + 2⋅c[i] + 3⋅d[i] = b[i + 1]" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Continuity of first derivatives at knots \n", + "dsi = diff(si, t)\n", + "eq3 = Eq(dsi.subs(t,1), dsi.subs(i, i+1).subs(t,0))\n", + "eq3" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$$2 {c}_{i} + 6 {d}_{i} = 2 {c}_{i + 1}$$" + ], + "text/plain": [ + "2⋅c[i] + 6⋅d[i] = 2⋅c[i + 1]" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Continuity of second derivatives at knots\n", + "d2si = diff(si, t, 2)\n", + "eq4 = Eq(d2si.subs(t,1), d2si.subs(i, i+1).subs(t,0))\n", + "eq4" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Boundary conditions" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$$2 {c}_{0} = 0$$" + ], + "text/plain": [ + "2⋅c[0] = 0" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/latex": [ + "$$2 {c}_{n - 1} + 6 {d}_{n - 1} = 0$$" + ], + "text/plain": [ + "2⋅c[n - 1] + 6⋅d[n - 1] = 0" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Natural BC (second derivative is zero at both ends)\n", + "eq5 = Eq(d2si.subs(i,0).subs(t,0), 0)\n", + "display(eq5)\n", + "eq6 = Eq(d2si.subs(i,n-1).subs(t,1), 0)\n", + "eq6" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$${b}_{0} = 0$$" + ], + "text/plain": [ + "b[0] = 0" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/latex": [ + "$${b}_{n - 1} + 2 {c}_{n - 1} + 3 {d}_{n - 1} = 0$$" + ], + "text/plain": [ + "b[n - 1] + 2⋅c[n - 1] + 3⋅d[n - 1] = 0" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Or\n", + "# Specified first derivatives at boundaries\n", + "eq5b = Eq(dsi.subs(i,0).subs(t,0),0)\n", + "display(eq5b)\n", + "eq6b = Eq(dsi.subs(i,n-1).subs(t,1),0)\n", + "eq6b" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Set up equations\n", + "\n", + "There are 4$n$ unknowns. To get the right number of equations, we need \n", + " - one of eq5 and eq6 (boundary conditions) (2 equations)\n", + " - eq1 and eq2 (values at beginning and end of each interval) for every knot ($2n$ equations)\n", + " - eq3 and eq4 (continuity of 1st and 2nd derivatives) for every interior knot ($2(n-1)$ equations)\n", + "\n", + "This gives the necessary $4n$ equations" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$${a}_{i} = {y}_{i}$$" + ], + "text/plain": [ + "a[i] = y[i]" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/latex": [ + "$${a}_{i} + {b}_{i} + {c}_{i} + {d}_{i} = {y}_{i + 1}$$" + ], + "text/plain": [ + "a[i] + b[i] + c[i] + d[i] = y[i + 1]" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/latex": [ + "$${b}_{i} + 2 {c}_{i} + 3 {d}_{i} = {b}_{i + 1}$$" + ], + "text/plain": [ + "b[i] + 2⋅c[i] + 3⋅d[i] = b[i + 1]" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/latex": [ + "$$2 {c}_{i} + 6 {d}_{i} = 2 {c}_{i + 1}$$" + ], + "text/plain": [ + "2⋅c[i] + 6⋅d[i] = 2⋅c[i + 1]" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "display(eq1)\n", + "display(eq2)\n", + "display(eq3)\n", + "display(eq4)" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$$- {b}_{i + 1} + {b}_{i} + 2 {c}_{i} + 3 {d}_{i} = 0$$" + ], + "text/plain": [ + "-b[i + 1] + b[i] + 2⋅c[i] + 3⋅d[i] = 0" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/latex": [ + "$$- 2 {c}_{i + 1} + 2 {c}_{i} + 6 {d}_{i} = 0$$" + ], + "text/plain": [ + "-2⋅c[i + 1] + 2⋅c[i] + 6⋅d[i] = 0" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "# Equations 3 and 4 need to be rearranged so all the unknowns are the left-hand side\n", + "sym_lhs = [a[i],b[i],c[i],d[i],a[i+1],b[i+1],c[i+1],d[i+1]]\n", + "sym_rhs = [y[i]]\n", + "eq3b = divide_terms(eq3, sym_lhs, sym_rhs)\n", + "display(eq3b)\n", + "eq4b = divide_terms(eq4, sym_lhs, sym_rhs)\n", + "display(eq4b)" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "metadata": {}, + "outputs": [], + "source": [ + "# Set up matrix with concrete number of knots\n", + "nval = 2\n", + "nknots = nval + 1\n", + "neqn = 4*nval\n", + "m = Matrix.eye(neqn, neqn)\n", + "rhs = Matrix.eye(neqn,1)\n", + "unknowns = Matrix.eye(neqn,1)" + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$$\\left[\\begin{matrix}0 & 0 & 2 & 0 & 0 & 0 & 0 & 0\\\\1 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\\\1 & 1 & 1 & 1 & 0 & 0 & 0 & 0\\\\0 & 0 & 0 & 0 & 1 & 0 & 0 & 0\\\\0 & 0 & 0 & 0 & 1 & 1 & 1 & 1\\\\0 & 1 & 2 & 3 & 0 & -1 & 0 & 0\\\\0 & 0 & 2 & 6 & 0 & 0 & -2 & 0\\\\0 & 0 & 0 & 0 & 0 & 0 & 2 & 6\\end{matrix}\\right] = \\left[\\begin{matrix}{a}_{0}\\\\{b}_{0}\\\\{c}_{0}\\\\{d}_{0}\\\\{a}_{1}\\\\{b}_{1}\\\\{c}_{1}\\\\{d}_{1}\\end{matrix}\\right] = \\left[\\begin{matrix}0\\\\{y}_{0}\\\\{y}_{1}\\\\{y}_{1}\\\\{y}_{2}\\\\0\\\\0\\\\0\\end{matrix}\\right]$$" + ], + "text/plain": [ + "⎡0 0 2 0 0 0 0 0⎤ ⎡a[0]⎤ ⎡ 0 ⎤\n", + "⎢ ⎥ ⎢ ⎥ ⎢ ⎥\n", + "⎢1 0 0 0 0 0 0 0⎥ ⎢b[0]⎥ ⎢y[0]⎥\n", + "⎢ ⎥ ⎢ ⎥ ⎢ ⎥\n", + "⎢1 1 1 1 0 0 0 0⎥ ⎢c[0]⎥ ⎢y[1]⎥\n", + "⎢ ⎥ ⎢ ⎥ ⎢ ⎥\n", + "⎢0 0 0 0 1 0 0 0⎥ ⎢d[0]⎥ ⎢y[1]⎥\n", + "⎢ ⎥ = ⎢ ⎥ = ⎢ ⎥\n", + "⎢0 0 0 0 1 1 1 1⎥ ⎢a[1]⎥ ⎢y[2]⎥\n", + "⎢ ⎥ ⎢ ⎥ ⎢ ⎥\n", + "⎢0 1 2 3 0 -1 0 0⎥ ⎢b[1]⎥ ⎢ 0 ⎥\n", + "⎢ ⎥ ⎢ ⎥ ⎢ ⎥\n", + "⎢0 0 2 6 0 0 -2 0⎥ ⎢c[1]⎥ ⎢ 0 ⎥\n", + "⎢ ⎥ ⎢ ⎥ ⎢ ⎥\n", + "⎣0 0 0 0 0 0 2 6⎦ ⎣d[1]⎦ ⎣ 0 ⎦" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "use_natural_bc = True\n", + "symlist = list()\n", + "for idx in range(nval):\n", + " symlist.append(a[idx])\n", + " symlist.append(b[idx])\n", + " symlist.append(c[idx])\n", + " symlist.append(d[idx])\n", + " \n", + "for idx,sym in enumerate(symlist):\n", + " unknowns[idx] = sym\n", + "\n", + "# First rows\n", + "for idx,sym in enumerate(symlist):\n", + " if use_natural_bc: \n", + " m[0,idx] = get_coeff_for(eq5.lhs.subs(i,0), sym , symlist)\n", + " else:\n", + " m[0,idx] = get_coeff_for(eq5b.lhs.subs(i,0), sym, symlist)\n", + " m[1,idx] = get_coeff_for(eq1.lhs.subs(i,0), sym,symlist)\n", + " m[2,idx] = get_coeff_for(eq2.lhs.subs(i,0), sym,symlist)\n", + " \n", + "rhs[0] = eq5.rhs.subs(i,0)\n", + "rhs[1] = eq1.rhs.subs(i,0)\n", + "rhs[2] = eq2.rhs.subs(i,0)\n", + "\n", + "jdx = 3\n", + "for nv in range(1,nval):\n", + " for idx,sym in enumerate(symlist):\n", + " m[jdx+0,idx] = get_coeff_for(eq1.lhs.subs(i,nv), sym,symlist)\n", + " m[jdx+1,idx] = get_coeff_for(eq2.lhs.subs(i,nv), sym,symlist)\n", + " m[jdx+2,idx] = get_coeff_for(eq3b.lhs.subs(i,nv-1), sym,symlist)\n", + " m[jdx+3,idx] = get_coeff_for(eq4b.lhs.subs(i,nv-1), sym,symlist)\n", + " rhs[jdx+0] = eq1.rhs.subs(i,nv)\n", + " rhs[jdx+1] = eq2.rhs.subs(i,nv)\n", + " rhs[jdx+2] = eq3b.rhs.subs(i,nv)\n", + " rhs[jdx+3] = eq4b.rhs.subs(i,nv)\n", + " jdx += 4\n", + "\n", + "# Last rows\n", + "for idx,sym in enumerate(symlist):\n", + " if use_natural_bc:\n", + " m[jdx,idx] = get_coeff_for(eq6.lhs.subs(n,nval), sym , symlist)\n", + " else:\n", + " m[jdx,idx] = get_coeff_for(eq6b.lhs.subs(n,nval), sym , symlist)\n", + "rhs[jdx] = eq5.rhs.subs(n,nval)\n", + "\n", + "# Display of matrix-vector multiplication is not working - pretend the first equals is a multiplication\n", + "with evaluate(False):\n", + " display(Eq(Eq(m,unknowns),rhs))" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$$\\left[\\begin{matrix}{y}_{0}\\\\- \\frac{5 {y}_{0}}{4} + \\frac{3 {y}_{1}}{2} - \\frac{{y}_{2}}{4}\\\\0\\\\\\frac{{y}_{0}}{4} - \\frac{{y}_{1}}{2} + \\frac{{y}_{2}}{4}\\\\{y}_{1}\\\\- \\frac{{y}_{0}}{2} + \\frac{{y}_{2}}{2}\\\\\\frac{3 {y}_{0}}{4} - \\frac{3 {y}_{1}}{2} + \\frac{3 {y}_{2}}{4}\\\\- \\frac{{y}_{0}}{4} + \\frac{{y}_{1}}{2} - \\frac{{y}_{2}}{4}\\end{matrix}\\right]$$" + ], + "text/plain": [ + "⎡ y[0] ⎤\n", + "⎢ ⎥\n", + "⎢ 5⋅y[0] 3⋅y[1] y[2]⎥\n", + "⎢- ────── + ────── - ────⎥\n", + "⎢ 4 2 4 ⎥\n", + "⎢ ⎥\n", + "⎢ 0 ⎥\n", + "⎢ ⎥\n", + "⎢ y[0] y[1] y[2] ⎥\n", + "⎢ ──── - ──── + ──── ⎥\n", + "⎢ 4 2 4 ⎥\n", + "⎢ ⎥\n", + "⎢ y[1] ⎥\n", + "⎢ ⎥\n", + "⎢ y[0] y[2] ⎥\n", + "⎢ - ──── + ──── ⎥\n", + "⎢ 2 2 ⎥\n", + "⎢ ⎥\n", + "⎢3⋅y[0] 3⋅y[1] 3⋅y[2]⎥\n", + "⎢────── - ────── + ──────⎥\n", + "⎢ 4 2 4 ⎥\n", + "⎢ ⎥\n", + "⎢ y[0] y[1] y[2] ⎥\n", + "⎢ - ──── + ──── - ──── ⎥\n", + "⎣ 4 2 4 ⎦" + ] + }, + "execution_count": 23, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Solve\n", + "sln = m.LUsolve(rhs)\n", + "sln" + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$$\\left \\{ {y}_{0} : 0.5, \\quad {y}_{1} : 1.1, \\quad {y}_{2} : 1.5\\right \\}$$" + ], + "text/plain": [ + "{y[0]: 0.5, y[1]: 1.1, y[2]: 1.5}" + ] + }, + "execution_count": 25, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Can substitute concrete knot values\n", + "ysubs = {y[0]:0.5, y[1]:1.1, y[2]:1.5}\n", + "ysubs" + ] + }, + { + "cell_type": "code", + "execution_count": 36, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$$\\left \\{ {a}_{0} : {y}_{0}, \\quad {a}_{1} : {y}_{1}, \\quad {b}_{0} : - \\frac{5 {y}_{0}}{4} + \\frac{3 {y}_{1}}{2} - \\frac{{y}_{2}}{4}, \\quad {b}_{1} : - \\frac{{y}_{0}}{2} + \\frac{{y}_{2}}{2}, \\quad {c}_{0} : 0, \\quad {c}_{1} : \\frac{3 {y}_{0}}{4} - \\frac{3 {y}_{1}}{2} + \\frac{3 {y}_{2}}{4}, \\quad {d}_{0} : \\frac{{y}_{0}}{4} - \\frac{{y}_{1}}{2} + \\frac{{y}_{2}}{4}, \\quad {d}_{1} : - \\frac{{y}_{0}}{4} + \\frac{{y}_{1}}{2} - \\frac{{y}_{2}}{4}\\right \\}$$" + ], + "text/plain": [ + "⎧ 5⋅y[0] 3⋅y[1] y[2] y[0] y[2] \n", + "⎨a[0]: y[0], a[1]: y[1], b[0]: - ────── + ────── - ────, b[1]: - ──── + ────, \n", + "⎩ 4 2 4 2 2 \n", + "\n", + " 3⋅y[0] 3⋅y[1] 3⋅y[2] y[0] y[1] y[2] y[0\n", + "c[0]: 0, c[1]: ────── - ────── + ──────, d[0]: ──── - ──── + ────, d[1]: - ───\n", + " 4 2 4 4 2 4 4 \n", + "\n", + "] y[1] y[2]⎫\n", + "─ + ──── - ────⎬\n", + " 2 4 ⎭" + ] + }, + "execution_count": 36, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Create substitutions for spline coefficients\n", + "subslist = dict(zip(symlist,sln))\n", + "subslist" + ] + }, + { + "cell_type": "code", + "execution_count": 40, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$$t^{3} \\left(\\frac{{y}_{0}}{4} - \\frac{{y}_{1}}{2} + \\frac{{y}_{2}}{4}\\right) + t \\left(- \\frac{5 {y}_{0}}{4} + \\frac{3 {y}_{1}}{2} - \\frac{{y}_{2}}{4}\\right) + {y}_{0}$$" + ], + "text/plain": [ + " 3 ⎛y[0] y[1] y[2]⎞ ⎛ 5⋅y[0] 3⋅y[1] y[2]⎞ \n", + "t ⋅⎜──── - ──── + ────⎟ + t⋅⎜- ────── + ────── - ────⎟ + y[0]\n", + " ⎝ 4 2 4 ⎠ ⎝ 4 2 4 ⎠ " + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/latex": [ + "$$- 0.05 t^{3} + 0.65 t + 0.5$$" + ], + "text/plain": [ + " 3 \n", + "- 0.05⋅t + 0.65⋅t + 0.5" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "# For a particular interval\n", + "si0 = si.subs(i,0).subs(subslist)\n", + "display(si0)\n", + "si0b = si0.subs(ysubs)\n", + "display(si0b)" + ] + }, + { + "cell_type": "code", + "execution_count": 45, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$$0$$" + ], + "text/plain": [ + "0" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/latex": [ + "$$-0.3$$" + ], + "text/plain": [ + "-0.300000000000000" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "# Can take the seconnd derivative, to compare with the output values from some spline solvers\n", + "# At t=0, values is zero, as expected from natural boundary conditions\n", + "display(diff(si0b, t, 2).subs(t,0))\n", + "# Compute at t=1\n", + "display(diff(si0b, t, 2).subs(t,1))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Derive for general knot spacing" + ] + }, + { + "cell_type": "code", + "execution_count": 51, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$${y}_{i} = {a}_{i} + {b}_{i} {t}_{i} + {c}_{i} {t}_{i}^{2} + {d}_{i} {t}_{i}^{3}$$" + ], + "text/plain": [ + " 2 3\n", + "y[i] = a[i] + b[i]⋅t[i] + c[i]⋅t[i] + d[i]⋅t[i] " + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "# Distance from the previous knot\n", + "# t[i] ranges from 0 to x[i+1]-x[i]\n", + "t = IndexedBase('t')\n", + "\n", + "# Number of knots\n", + "n = Symbol('n', integer=True)\n", + "i = Symbol('i', integer=True)\n", + "# Location of knots\n", + "x = IndexedBase('x')\n", + "# Function values to intepolated at the knots\n", + "y = IndexedBase('y')\n", + "\n", + "# Coefficients of the spline function\n", + "a,b,c,d = [IndexedBase(s) for s in 'a b c d'.split()]\n", + "\n", + "\n", + "# With indexed variables\n", + "si = a[i] + b[i]*t[i] + c[i]*t[i]*t[i] + d[i]*t[i]**3\n", + "display(Eq(y[i],si))" + ] + }, + { + "cell_type": "code", + "execution_count": 50, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$${a}_{i} = {y}_{i}$$" + ], + "text/plain": [ + "a[i] = y[i]" + ] + }, + "execution_count": 50, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Value at knots (t[i]=0)\n", + "eq1 = Eq(si.subs(t[i],0), y[i])\n", + "eq1" + ] + }, + { + "cell_type": "code", + "execution_count": 55, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$$\\left({x}_{i + 1} - {x}_{i}\\right)^{3} {d}_{i} + \\left({x}_{i + 1} - {x}_{i}\\right)^{2} {c}_{i} + \\left({x}_{i + 1} - {x}_{i}\\right) {b}_{i} + {a}_{i} = {y}_{i + 1}$$" + ], + "text/plain": [ + " 3 2 \n", + "(x[i + 1] - x[i]) ⋅d[i] + (x[i + 1] - x[i]) ⋅c[i] + (x[i + 1] - x[i])⋅b[i] + a\n", + "\n", + " \n", + "[i] = y[i + 1]" + ] + }, + "execution_count": 55, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Value at knots (t[i]=x[i+1])\n", + "eq2 = Eq(si.subs(t[i],x[i+1]-x[i]), y[i+1])\n", + "eq2" + ] + }, + { + "cell_type": "code", + "execution_count": 96, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$$3 \\left({x}_{i + 1} - {x}_{i}\\right)^{2} {d}_{i} + 2 \\left({x}_{i + 1} - {x}_{i}\\right) {c}_{i} + {b}_{i} = {b}_{i + 1}$$" + ], + "text/plain": [ + " 2 \n", + "3⋅(x[i + 1] - x[i]) ⋅d[i] + 2⋅(x[i + 1] - x[i])⋅c[i] + b[i] = b[i + 1]" + ] + }, + "execution_count": 96, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Continuity of first derivatives at knots \n", + "dsi = diff(si, t[i])\n", + "eq3 = Eq(dsi.subs(t[i],x[i+1]-x[i]), dsi.subs(i, i+1).subs(t[i+1],0))\n", + "eq3" + ] + }, + { + "cell_type": "code", + "execution_count": 97, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$$6 \\left({x}_{i + 1} - {x}_{i}\\right) {d}_{i} + 2 {c}_{i} = 2 {c}_{i + 1}$$" + ], + "text/plain": [ + "6⋅(x[i + 1] - x[i])⋅d[i] + 2⋅c[i] = 2⋅c[i + 1]" + ] + }, + "execution_count": 97, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Continuity of second derivatives at knots\n", + "d2si = diff(si, t[i], 2)\n", + "eq4 = Eq(d2si.subs(t[i],x[i+1]-x[i]), d2si.subs(i, i+1).subs(t[i+1],0))\n", + "eq4" + ] + }, + { + "cell_type": "code", + "execution_count": 118, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$$2 {c}_{0} = 0$$" + ], + "text/plain": [ + "2⋅c[0] = 0" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/latex": [ + "$$6 \\left(- {x}_{n - 1} + {x}_{n}\\right) {d}_{n - 1} + 2 {c}_{n - 1} = 0$$" + ], + "text/plain": [ + "6⋅(-x[n - 1] + x[n])⋅d[n - 1] + 2⋅c[n - 1] = 0" + ] + }, + "execution_count": 118, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Natural BC (second derivative is zero at both ends)\n", + "eq5 = Eq(d2si.subs(i,0).subs(t[0],0), 0)\n", + "display(eq5)\n", + "eq6 = Eq(d2si.subs(t[i],x[i+1]-x[i]).subs(i,n-1), 0)\n", + "eq6" + ] + }, + { + "cell_type": "code", + "execution_count": 117, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$${b}_{0} = 0$$" + ], + "text/plain": [ + "b[0] = 0" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/latex": [ + "$$3 \\left(- {x}_{n - 1} + {x}_{n}\\right)^{2} {d}_{n - 1} + 2 \\left(- {x}_{n - 1} + {x}_{n}\\right) {c}_{n - 1} + {b}_{n - 1} = 0$$" + ], + "text/plain": [ + " 2 \n", + "3⋅(-x[n - 1] + x[n]) ⋅d[n - 1] + 2⋅(-x[n - 1] + x[n])⋅c[n - 1] + b[n - 1] = 0" + ] + }, + "execution_count": 117, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Or\n", + "# Specified first derivatives at boundaries\n", + "eq5b = Eq(dsi.subs(i,0).subs(t[0],0),0)\n", + "display(eq5b)\n", + "eq6b = Eq(dsi.subs(t[i],x[i+1]-x[i]).subs(i,n-1),0)\n", + "eq6b" + ] + }, + { + "cell_type": "code", + "execution_count": 99, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$$3 \\left({x}_{i + 1} - {x}_{i}\\right)^{2} {d}_{i} + 2 \\left({x}_{i + 1} - {x}_{i}\\right) {c}_{i} - {b}_{i + 1} + {b}_{i} = 0$$" + ], + "text/plain": [ + " 2 \n", + "3⋅(x[i + 1] - x[i]) ⋅d[i] + 2⋅(x[i + 1] - x[i])⋅c[i] - b[i + 1] + b[i] = 0" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/latex": [ + "$$6 \\left({x}_{i + 1} - {x}_{i}\\right) {d}_{i} - 2 {c}_{i + 1} + 2 {c}_{i} = 0$$" + ], + "text/plain": [ + "6⋅(x[i + 1] - x[i])⋅d[i] - 2⋅c[i + 1] + 2⋅c[i] = 0" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "# Equations 3 and 4 need to be rearranged so all the unknowns are the left-hand side\n", + "sym_lhs = [a[i],b[i],c[i],d[i],a[i+1],b[i+1],c[i+1],d[i+1]]\n", + "sym_rhs = [y[i]]\n", + "eq3b = divide_terms(eq3, sym_lhs, sym_rhs)\n", + "display(eq3b)\n", + "eq4b = divide_terms(eq4, sym_lhs, sym_rhs)\n", + "display(eq4b)" + ] + }, + { + "cell_type": "code", + "execution_count": 100, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$${a}_{i} = {y}_{i}$$" + ], + "text/plain": [ + "a[i] = y[i]" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/latex": [ + "$$\\left({x}_{i + 1} - {x}_{i}\\right)^{3} {d}_{i} + \\left({x}_{i + 1} - {x}_{i}\\right)^{2} {c}_{i} + \\left({x}_{i + 1} - {x}_{i}\\right) {b}_{i} + {a}_{i} = {y}_{i + 1}$$" + ], + "text/plain": [ + " 3 2 \n", + "(x[i + 1] - x[i]) ⋅d[i] + (x[i + 1] - x[i]) ⋅c[i] + (x[i + 1] - x[i])⋅b[i] + a\n", + "\n", + " \n", + "[i] = y[i + 1]" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/latex": [ + "$$3 \\left({x}_{i + 1} - {x}_{i}\\right)^{2} {d}_{i} + 2 \\left({x}_{i + 1} - {x}_{i}\\right) {c}_{i} - {b}_{i + 1} + {b}_{i} = 0$$" + ], + "text/plain": [ + " 2 \n", + "3⋅(x[i + 1] - x[i]) ⋅d[i] + 2⋅(x[i + 1] - x[i])⋅c[i] - b[i + 1] + b[i] = 0" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/latex": [ + "$$6 \\left({x}_{i + 1} - {x}_{i}\\right) {d}_{i} - 2 {c}_{i + 1} + 2 {c}_{i} = 0$$" + ], + "text/plain": [ + "6⋅(x[i + 1] - x[i])⋅d[i] - 2⋅c[i + 1] + 2⋅c[i] = 0" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "display(eq1)\n", + "display(eq2)\n", + "display(eq3b)\n", + "display(eq4b)" + ] + }, + { + "cell_type": "code", + "execution_count": 66, + "metadata": {}, + "outputs": [], + "source": [ + "# Set up matrix with concrete number of knots\n", + "nval = 2\n", + "nknots = nval + 1\n", + "neqn = 4*nval\n", + "m = Matrix.eye(neqn, neqn)\n", + "rhs = Matrix.eye(neqn,1)\n", + "unknowns = Matrix.eye(neqn,1)" + ] + }, + { + "cell_type": "code", + "execution_count": 116, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$$\\left[\\begin{matrix}0 & 0 & 2 & 0 & 0 & 0 & 0 & 0\\\\1 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\\\1 & - {x}_{0} + {x}_{1} & \\left(- {x}_{0} + {x}_{1}\\right)^{2} & \\left(- {x}_{0} + {x}_{1}\\right)^{3} & 0 & 0 & 0 & 0\\\\0 & 0 & 0 & 0 & 1 & 0 & 0 & 0\\\\0 & 0 & 0 & 0 & 1 & - {x}_{1} + {x}_{2} & \\left(- {x}_{1} + {x}_{2}\\right)^{2} & \\left(- {x}_{1} + {x}_{2}\\right)^{3}\\\\0 & 1 & - 2 {x}_{0} + 2 {x}_{1} & 3 \\left(- {x}_{0} + {x}_{1}\\right)^{2} & 0 & -1 & 0 & 0\\\\0 & 0 & 2 & - 6 {x}_{0} + 6 {x}_{1} & 0 & 0 & -2 & 0\\\\0 & 0 & 0 & 0 & 0 & 0 & 2 & - 6 {x}_{1} + 6 {x}_{2}\\end{matrix}\\right]$$" + ], + "text/plain": [ + "⎡0 0 2 0 0 0 \n", + "⎢ \n", + "⎢1 0 0 0 0 0 \n", + "⎢ \n", + "⎢ 2 3 \n", + "⎢1 -x[0] + x[1] (-x[0] + x[1]) (-x[0] + x[1]) 0 0 \n", + "⎢ \n", + "⎢0 0 0 0 1 0 \n", + "⎢ \n", + "⎢ \n", + "⎢0 0 0 0 1 -x[1] + x[2] (-x[1]\n", + "⎢ \n", + "⎢ 2 \n", + "⎢0 1 -2⋅x[0] + 2⋅x[1] 3⋅(-x[0] + x[1]) 0 -1 \n", + "⎢ \n", + "⎢0 0 2 -6⋅x[0] + 6⋅x[1] 0 0 \n", + "⎢ \n", + "⎣0 0 0 0 0 0 \n", + "\n", + " 0 0 ⎤\n", + " ⎥\n", + " 0 0 ⎥\n", + " ⎥\n", + " ⎥\n", + " 0 0 ⎥\n", + " ⎥\n", + " 0 0 ⎥\n", + " ⎥\n", + " 2 3 ⎥\n", + " + x[2]) (-x[1] + x[2]) ⎥\n", + " ⎥\n", + " ⎥\n", + " 0 0 ⎥\n", + " ⎥\n", + "-2 0 ⎥\n", + " ⎥\n", + " 2 -6⋅x[1] + 6⋅x[2]⎦" + ] + }, + "execution_count": 116, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "use_natural_bc = True\n", + "symlist = list()\n", + "for idx in range(nval):\n", + " symlist.append(a[idx])\n", + " symlist.append(b[idx])\n", + " symlist.append(c[idx])\n", + " symlist.append(d[idx])\n", + " \n", + "for idx,sym in enumerate(symlist):\n", + " unknowns[idx] = sym\n", + "\n", + "# First rows\n", + "for idx,sym in enumerate(symlist):\n", + " if use_natural_bc: \n", + " m[0,idx] = get_coeff_for(eq5.lhs.subs(i,0), sym , symlist)\n", + " else:\n", + " m[0,idx] = get_coeff_for(eq5b.lhs.subs(i,0), sym, symlist)\n", + " m[1,idx] = get_coeff_for(eq1.lhs.subs(i,0), sym,symlist)\n", + " m[2,idx] = get_coeff_for(eq2.lhs.subs(i,0), sym,symlist)\n", + " \n", + "rhs[0] = eq5.rhs.subs(i,0)\n", + "rhs[1] = eq1.rhs.subs(i,0)\n", + "rhs[2] = eq2.rhs.subs(i,0)\n", + "\n", + "jdx = 3\n", + "for nv in range(1,nval):\n", + " for idx,sym in enumerate(symlist):\n", + " m[jdx+0,idx] = get_coeff_for(eq1.lhs.subs(i,nv), sym,symlist)\n", + " m[jdx+1,idx] = get_coeff_for(eq2.lhs.subs(i,nv), sym,symlist)\n", + " m[jdx+2,idx] = get_coeff_for(eq3b.lhs.subs(i,nv-1), sym,symlist)\n", + " m[jdx+3,idx] = get_coeff_for(eq4b.lhs.subs(i,nv-1), sym,symlist)\n", + " rhs[jdx+0] = eq1.rhs.subs(i,nv)\n", + " rhs[jdx+1] = eq2.rhs.subs(i,nv)\n", + " rhs[jdx+2] = eq3b.rhs.subs(i,nv)\n", + " rhs[jdx+3] = eq4b.rhs.subs(i,nv)\n", + " jdx += 4\n", + "\n", + "# Last rows\n", + "for idx,sym in enumerate(symlist):\n", + " if use_natural_bc:\n", + " m[jdx,idx] = get_coeff_for(eq6.lhs.subs(n,nval), sym , symlist)\n", + " else:\n", + " m[jdx,idx] = get_coeff_for(eq6b.lhs.subs(n,nval), sym , symlist)\n", + "rhs[jdx] = eq5.rhs.subs(n,nval)\n", + "\n", + "# Display just the matrix, displaying other pieces doesn't work, and they are the same as previously\n", + "m" + ] + }, + { + "cell_type": "code", + "execution_count": 103, + "metadata": {}, + "outputs": [], + "source": [ + "# Solve\n", + "sln = m.LUsolve(rhs)\n", + "# Display is large - skip for now, uncomment if you want to see it\n", + "#sln" + ] + }, + { + "cell_type": "code", + "execution_count": 104, + "metadata": {}, + "outputs": [], + "source": [ + "# Create substitutions for spline coefficients\n", + "subslist = dict(zip(symlist,sln))\n", + "#subslist" + ] + }, + { + "cell_type": "code", + "execution_count": 105, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$$\\left \\{ {y}_{0} : 0.5, \\quad {y}_{1} : 1.1, \\quad {y}_{2} : 1.5\\right \\}$$" + ], + "text/plain": [ + "{y[0]: 0.5, y[1]: 1.1, y[2]: 1.5}" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/latex": [ + "$$\\left \\{ {x}_{0} : 0.0, \\quad {x}_{1} : 1.0, \\quad {x}_{2} : 2.0\\right \\}$$" + ], + "text/plain": [ + "{x[0]: 0.0, x[1]: 1.0, x[2]: 2.0}" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "# Can substitute concrete knot values\n", + "ysubs = {y[0]:0.5, y[1]:1.1, y[2]:1.5}\n", + "display(ysubs)\n", + "# Use uniform knots to verify answer is the same as before\n", + "xsubs = {x[0]:0.0, x[1]:1.0, x[2]:2.0}\n", + "display(xsubs)" + ] + }, + { + "cell_type": "code", + "execution_count": 112, + "metadata": {}, + "outputs": [ + { + "data": { + "text/latex": [ + "$$- 0.05 t^{3} + 0.65 t + 0.5$$" + ], + "text/plain": [ + " 3 \n", + "- 0.05⋅t + 0.65⋅t + 0.5" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "# For a particular interval\n", + "si0 = si.subs(i,0).subs(subslist)\n", + "#display(si0)\n", + "si0b = si0.subs(ysubs).subs(xsubs).subs(t[0],t)\n", + "display(si0b)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.7.1" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/Wavefunctions/eqn_manip.py b/Wavefunctions/eqn_manip.py new file mode 100644 index 0000000..e167202 --- /dev/null +++ b/Wavefunctions/eqn_manip.py @@ -0,0 +1,49 @@ +from __future__ import print_function + +# Code for manipulating equations (moving terms from one side to another, etc.) + +from sympy import Eq + +# Move symbols in sym_list from left hand side of equation to right hand side +def move_terms(eqn, sym_list): + new_lhs = eqn.lhs + new_rhs = eqn.rhs + for sym in sym_list: + c = eqn.lhs.coeff(sym) + new_lhs = new_lhs - c*sym + new_rhs = new_rhs - c*sym + return Eq(new_lhs, new_rhs) + +# Move symbols in sym_list from right hand side of equation to left hand side +def move_terms_left(eqn, sym_list): + new_lhs = eqn.lhs + new_rhs = eqn.rhs + for sym in sym_list: + c = eqn.rhs.coeff(sym) + new_lhs = new_lhs - c*sym + new_rhs = new_rhs - c*sym + return Eq(new_lhs, new_rhs) + +def divide_terms(eqn, sym_list_left, sym_list_right): + #print 'start',eqn + eqn1 = move_terms(eqn, sym_list_right) + #print 'middle ',eqn1 + eqn2 = move_terms_left(eqn1, sym_list_left) + return eqn2 + +# Multiply equation by term +def mult_eqn(eqn, e): + return Eq(eqn.lhs*e, eqn.rhs*e) + +# for all values other than the target, set to zero +def get_coeff_for(expr, sym, symlist): + # expression + # symbol to get coefficient for + # symlist - total list of symbols + subslist = {} + subslist[sym] = 1 + for s in symlist: + if s != sym: + subslist[s] = 0 + coeff = expr.subs(subslist) + return coeff