├── display_nnout.py ├── save_nn.py ├── README.md ├── schroedinger_nn.py ├── visualize_nn.py └── genpotential.py /display_nnout.py: -------------------------------------------------------------------------------- 1 | display_nnout.py 2 | # Makes plots of an individual potential (scaled to unit max), the gradient descent (“correct”) ground state, 3 | # and the neural network predicted ground state 4 | # should be added to notebook containing schroedinger_nn.py 5 | import matplotlib.pyplot as mp 6 | potenid = 5397 7 | mp.plot(sess.run(L3,feed_dict={X: [trainx[potenid]]})[0]) 8 | mp.plot([trainx[potenid][i]/max(trainx[potenid]) for i in range(bins - 1)]) 9 | mp.plot(trainy[potenid]) 10 | mp.show() 11 | -------------------------------------------------------------------------------- /save_nn.py: -------------------------------------------------------------------------------- 1 | # save_nn.py 2 | # small tool to save the neural network state. append to schroedinger_nn.py notebook. 3 | import csv 4 | with open('W1.csv', 'w') as f: 5 | fileout = csv.writer(f) 6 | fileout.writerows(sess.run(W1).tolist()) 7 | with open('W2.csv', 'w') as f: 8 | fileout = csv.writer(f) 9 | fileout.writerows(sess.run(W2).tolist()) 10 | with open('W3.csv', 'w') as f: 11 | fileout = csv.writer(f) 12 | fileout.writerows(sess.run(W3).tolist()) 13 | with open('B1.csv', 'w') as f: 14 | fileout = csv.writer(f) 15 | fileout.writerows([sess.run(B1).tolist()]) 16 | with open('B2.csv', 'w') as f: 17 | fileout = csv.writer(f) 18 | fileout.writerows([sess.run(B2).tolist()]) 19 | with open('B3.csv', 'w') as f: 20 | fileout = csv.writer(f) 21 | fileout.writerows([sess.run(B3).tolist()]) 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # README.md 2 | # Schroedinger equation deep learning project 3 | # Version 1: 25 September 2017 4 | # Steven K. Steinke 5 | # steven.k.steinke@gmail.com 6 | This project uses Tensorflow to first generate random 1-D potentials, then solve them using 7 | gradient descent. These potentials and solutions are partitioned into training and validation data. 8 | Next, the training data are fed into a simple neural network with 2 hidden layers that use 9 | the softplus activation function. The mean squared distance between the “correct” solutions 10 | and the output of the neural network is the cost function; gradient descent on the network 11 | “solves” the problem. There are various simple tools included to visualize the output of this process. 12 | 13 | # Files included: 14 | genpotential.py Generates the random potentials and solves them individually 15 | 16 | schroedinger_nn.py Sets up a simple neural network to solve the 1D Sch. Eqn. 17 | 18 | display_nnout.py Plots a single potential and its actual and predicted solutions 19 | 20 | save_nn.py Saves the weights and biases of the network for later recovery 21 | 22 | visualize_nn.py Creates bitmaps of the weights and biases of the network after sorting for spatial correlation 23 | -------------------------------------------------------------------------------- /schroedinger_nn.py: -------------------------------------------------------------------------------- 1 | # schroedinger_nn.py 2 | # This reads the potential training data from genpotential.py and then sets up a neural network with 2 hidden layers. 3 | # Additional tools to output visualize and save the network are in other files. 4 | import csv 5 | import numpy as np 6 | import tensorflow as tf 7 | 8 | bins = 128 9 | seedmax = 20 # opens seed files 0 - 19. Lost too much data due to kernel crashes, so these got broken up 10 | trainx = [] 11 | trainy = [] 12 | validx = [] 13 | validy = [] 14 | 15 | #This is not a ... pythonic [barf]... way of reading data, but python is stupid about pointers, so deal with it 16 | for i in range(seedmax): 17 | with open('test_pots'+str(i)+'.csv', 'r') as csvfile: 18 | flurg = csv.reader(csvfile) 19 | for row in flurg: 20 | trainx.append([float(num) for num in row]) 21 | with open('test_out'+str(i)+'.csv', 'r') as csvfile: 22 | flurg = csv.reader(csvfile) 23 | for row in flurg: 24 | trainy.append([float(num) for num in row]) 25 | with open('valid_pots'+str(i)+'.csv', 'r') as csvfile: 26 | flurg = csv.reader(csvfile) 27 | for row in flurg: 28 | validx.append([float(num) for num in row]) 29 | with open('valid_out'+str(i)+'.csv', 'r') as csvfile: 30 | flurg = csv.reader(csvfile) 31 | for row in flurg: 32 | validy.append([float(num) for num in row]) 33 | 34 | seed = 42 35 | np.random.seed(seed) 36 | tf.set_random_seed(seed) 37 | #have a decaying learning rate so that convergence is faster at first and the fit is better at the end. 38 | #However, by trial and error, the simple exponential decay doesn't work well. 39 | #Trying a method by which the decay happens at hand-specified intervals 40 | startrate = 0.125 41 | gs = 0 42 | gslist = [1,1,2,3,10,20,40,100,200,10000] 43 | ic = 0 44 | learnrate = tf.Variable(startrate, trainable=False) 45 | updatelearnrate = tf.assign(learnrate,tf.multiply(learnrate,0.75)) 46 | # set up neural network layers. There are shorter ways to do it, but this exposes the guts. 47 | X = tf.placeholder(tf.float32) 48 | Y = tf.placeholder(tf.float32) 49 | #1st hidden layer 50 | W1 = tf.Variable(tf.random_uniform([bins-1, bins-1], -1./bins, 1./bins)) 51 | B1 = tf.Variable(tf.random_uniform([bins-1], -1., 1.)) 52 | L1 = tf.nn.softplus(tf.matmul(X, W1) + B1) 53 | #2nd hidden layer 54 | W2 = tf.Variable(tf.random_uniform([bins-1, bins-1], -1./bins, 1./bins)) 55 | B2 = tf.Variable(tf.random_uniform([bins-1], -1., 1.)) 56 | L2 = tf.nn.softplus(tf.matmul(L1, W2) + B2) 57 | #Output layer 58 | W3 = tf.Variable(tf.random_uniform([bins-1, bins-1], -1./bins, 1./bins)) 59 | B3 = tf.Variable(tf.random_uniform([bins-1], -1., 1.)) 60 | L3 = tf.nn.softplus(tf.matmul(L2, W3) + B3) 61 | #Cost function 62 | costfunc = tf.reduce_mean(tf.square(tf.subtract(L3,Y))) 63 | optimizer = tf.train.GradientDescentOptimizer(learnrate) 64 | trainstep = optimizer.minimize(costfunc) 65 | #initialize 66 | init = tf.global_variables_initializer() 67 | sess = tf.Session() 68 | sess.run(init) 69 | 70 | for step in range(100000): 71 | if step % 150 == 0: 72 | if ic == gslist[gs]: 73 | gs = gs + 1 74 | ic = 1 75 | sess.run(updatelearnrate) 76 | else: 77 | ic = ic + 1 78 | if step %100 == 0: 79 | print step, 'Train loss: ',sess.run(costfunc,feed_dict={X: trainx, Y: trainy}), 'Valid loss: ',sess.run(costfunc,feed_dict={X: validx, Y: validy}) 80 | sess.run(trainstep, feed_dict={X: trainx, Y: trainy}) 81 | 82 | 83 | -------------------------------------------------------------------------------- /visualize_nn.py: -------------------------------------------------------------------------------- 1 | # visualize_nn.py 2 | # outputs bitmaps of the weights and biases from schroedinger_nn.py. 3 | # Sorts them using a Gaussian kernel to increase spatial correlation between weights and nodes. 4 | # Doubles the size of the bitmap before outputting. 5 | # Append to schroedinger_nn.py notebook 6 | 7 | def doubler(aray): 8 | dbled = np.zeros([2*i for i in aray.shape]) 9 | if len(aray.shape) == 1: 10 | for i in range(aray.shape[0]): 11 | dbled[2*i] = aray[i] 12 | dbled[2*i+1] = aray[i] 13 | elif len(aray.shape) == 2: 14 | for i in range(aray.shape[0]): 15 | for j in range(aray.shape[1]): 16 | dbled[2*i][2*j] = aray[i][j] 17 | dbled[2*i+1][2*j] = aray[i][j] 18 | dbled[2*i][2*j+1] = aray[i][j] 19 | dbled[2*i+1][2*j+1] = aray[i][j] 20 | return dbled 21 | 22 | from PIL import Image 23 | we1 = sess.run(W1) 24 | bi1 = sess.run(B1) 25 | we2 = sess.run(W2) 26 | bi2 = sess.run(B2) 27 | we3 = sess.run(W3) 28 | bi3 = sess.run(B3) 29 | 30 | gauswid = 1. 31 | weiscale = [] 32 | for i in range(bins-1): 33 | line = np.exp([-np.square(float(i-j)/gauswid)/2. for j in range(bins-1)]) 34 | line = np.divide(line,sum(line)) 35 | weiscale.append(line.tolist()) 36 | 37 | weconv1 = np.matmul(we1,weiscale) 38 | weconv2 = np.matmul(we2,weiscale) 39 | 40 | sign = 1 41 | mask = np.zeros(bins-1) 42 | for i in range(bins-1): 43 | ind = (bins-2)/2+int(np.floor((i+1)/2))*sign 44 | sign = -sign 45 | mxin = np.argmax(np.add(weconv1[ind],mask)) 46 | swapper = np.identity(bins-1) 47 | swapper[ind][ind] = 0 48 | swapper[mxin][mxin] = 0 49 | swapper[ind][mxin] = 1 50 | swapper[mxin][ind] = 1 51 | we1 = np.matmul(we1,swapper) 52 | weconv1 = np.matmul(weconv1,swapper) 53 | bi1 = np.matmul(bi1,swapper) 54 | we2 = np.matmul(swapper,we2) 55 | mask[ind] = -1.E12 56 | 57 | sign = 1 58 | mask = np.zeros(bins-1) 59 | for i in range(bins-1): 60 | ind = (bins-2)/2+int(np.floor((i+1)/2))*sign 61 | sign = -sign 62 | mxin = np.argmax(np.add(weconv2[ind],mask)) 63 | swapper = np.identity(bins-1) 64 | swapper[ind][ind] = 0 65 | swapper[mxin][mxin] = 0 66 | swapper[ind][mxin] = 1 67 | swapper[mxin][ind] = 1 68 | we2 = np.matmul(we2,swapper) 69 | weconv2 = np.matmul(weconv2,swapper) 70 | bi2 = np.matmul(bi2,swapper) 71 | we3 = np.matmul(swapper,we3) 72 | mask[ind] = -1.E12 73 | 74 | 75 | max1 = max(max(we1.tolist())) 76 | min1 = min(min(we1.tolist())) 77 | wedb1 = doubler(we1) 78 | weight1 = np.divide(np.subtract(wedb1,min1),max1-min1) 79 | wim1 = Image.fromarray((weight1*255).astype(np.uint8),'L') 80 | wim1.save('W1.bmp') 81 | max1 = max(bi1.tolist()) 82 | min1 = min(bi1.tolist()) 83 | bidb1 = doubler(bi1) 84 | bia1 = np.divide(np.subtract(bidb1,min1),max1-min1) 85 | bias1 = np.array([bia1.tolist() for i in range(32)]) 86 | bim1 = Image.fromarray((bias1*255).astype(np.uint8),'L') 87 | bim1.save('B1.bmp') 88 | 89 | max2 = max(max(we2.tolist())) 90 | min2 = min(min(we2.tolist())) 91 | wedb2 = doubler(we2) 92 | weight2 = np.divide(np.subtract(wedb2,min2),max2-min2) 93 | wim2 = Image.fromarray((weight2*255).astype(np.uint8),'L') 94 | wim2.save('W2.bmp') 95 | max2 = max(bi2.tolist()) 96 | min2 = min(bi2.tolist()) 97 | bidb2 = doubler(bi2) 98 | bia2 = np.divide(np.subtract(bidb2,min2),max2-min2) 99 | bias2 = np.array([bia2.tolist() for i in range(32)]) 100 | bim2 = Image.fromarray((bias2*255).astype(np.uint8),'L') 101 | bim2.save('B2.bmp') 102 | 103 | max3 = max(max(we3.tolist())) 104 | min3 = min(min(we3.tolist())) 105 | wedb3 = doubler(we3) 106 | weight3 = np.divide(np.subtract(wedb3,min3),max3-min3) 107 | wim3 = Image.fromarray((weight3*255).astype(np.uint8),'L') 108 | wim3.save('W3.bmp') 109 | max3 = max(bi3.tolist()) 110 | min3 = min(bi3.tolist()) 111 | bidb3 = doubler(bi3) 112 | bia3 = np.divide(np.subtract(bidb3,min3),max3-min3) 113 | bias3 = np.array([bia3.tolist() for i in range(32)]) 114 | bim3 = Image.fromarray((bias3*255).astype(np.uint8),'L') 115 | bim3.save('B3.bmp') -------------------------------------------------------------------------------- /genpotential.py: -------------------------------------------------------------------------------- 1 | # genpotential.py 2 | # This snippet of python code generates random potentials of 3 different types, 3 | # Step functions, piecewise linear functions, and random Fourier series. 4 | # Each of these types gets more “jagged” as generation progresses. 5 | # The ground state wavefunction of each potential is found using Tensorflow’s gradient 6 | # descent method on the energy functional given by the Schroedinger equation. 7 | # The potentials and solutions are partitioned into training and validation data 8 | # and saved with the random seed appended to the filename 9 | 10 | import csv 11 | import numpy as np 12 | import tensorflow as tf 13 | 14 | def subexp(expon): 15 | return np.power(abs(np.log(np.random.uniform())),expon) 16 | 17 | def generatepot(style,param): #0=step,1=linear,2=fourier; 0-1 "jaggedness" scale 18 | mu = 1. + bins*param #mean number of jump points for styles 0 + 1 19 | forxp = 2.5 - 2*param #fourier exponent for style 2 20 | scale = 5.0*(np.pi*np.pi*0.5) # energy scale 21 | if style < 2: 22 | dx = bins/mu 23 | xlist = [-dx/2] 24 | while xlist[-1] < bins: 25 | xlist.append(xlist[-1]+dx*subexp(1.)) 26 | vlist = [scale*subexp(2.) for k in range(len(xlist))] 27 | k = 0 28 | poten = [] 29 | for l in range(1,bins): 30 | while xlist[k+1] < l: 31 | k = k + 1 32 | if style == 0: 33 | poten.append(vlist[k]) 34 | else: 35 | poten.append(vlist[k]+(vlist[k+1]-vlist[k])*(l-xlist[k])/(xlist[k+1]-xlist[k])) 36 | else: 37 | sincoef = [(2*np.random.randint(2)-1.)*scale*subexp(2.)/np.power(k,forxp) for k in range(1,bins//2)] 38 | coscoef = [(2*np.random.randint(2)-1.)*scale*subexp(2.)/np.power(k,forxp) for k in range(1,bins//2)] 39 | zercoef = scale*subexp(2.) 40 | poten = np.maximum(np.add(np.add(np.matmul(sincoef,sinval),np.matmul(coscoef,cosval)),zercoef),0).tolist() 41 | return poten 42 | 43 | seed = 0 44 | np.random.seed(seed) 45 | bins = 128 #dx = 1/bins; actual number of columns saved = bins-1, because 1st and last are 0 46 | npots = 200 #ends up being 3*this*(validnth-1)/validnth 47 | validnth = 5 #every nth sample func is saved as validation 48 | sinval = np.sin([[np.pi*i*j/bins for i in range(1,bins)] for j in range(1,bins//2)]) 49 | cosval = np.cos([[np.pi*i*j/bins for i in range(1,bins)] for j in range(1,bins//2)]) 50 | sqrt2 = np.sqrt(2) 51 | 52 | defgrdstate = tf.constant([sqrt2*np.sin(i*np.pi/bins) for i in range(1,bins)]) 53 | psi = tf.Variable(defgrdstate) 54 | zerotens = tf.zeros([1]) 55 | psil = tf.concat([psi[1:],zerotens],0) 56 | psir = tf.concat([zerotens,psi[:-1]],0) 57 | renorm = tf.assign(psi,tf.divide(psi,tf.sqrt(tf.reduce_mean(tf.square(psi))))) 58 | optimzi = tf.train.GradientDescentOptimizer(0.0625/bins) 59 | reinit = tf.assign(psi,defgrdstate) 60 | init = tf.global_variables_initializer() 61 | 62 | potentials = [] 63 | validpots = [] 64 | wavefuncs = [] 65 | validfuncs = [] 66 | 67 | sess = tf.Session() 68 | sess.run(init) 69 | for i in range(npots): 70 | if i%10 == 0: 71 | print str((100.*i)/npots) + '% complete' 72 | for j in range(3): 73 | vofx = generatepot(j,(1.*i)/npots) 74 | energy = tf.reduce_mean(tf.subtract(tf.multiply(tf.square(psi),tf.add(vofx,1.*bins*bins)), 75 | tf.multiply(tf.multiply(tf.add(psil,psir),psi),0.5*bins*bins))) 76 | training = optimzi.minimize(energy) 77 | sess.run(reinit) 78 | for t in range(20000): 79 | sess.run(training) 80 | sess.run(renorm) 81 | if i%validnth == 0: 82 | validpots.append(vofx) 83 | validfuncs.append(sess.run(psi).tolist()) 84 | else: 85 | potentials.append(vofx) 86 | wavefuncs.append(sess.run(psi).tolist()) 87 | 88 | with open('test_pots'+str(seed)+'.csv', 'w') as f: 89 | fileout = csv.writer(f) 90 | fileout.writerows(potentials) 91 | with open('valid_pots'+str(seed)+'.csv', 'w') as f: 92 | fileout = csv.writer(f) 93 | fileout.writerows(validpots) 94 | with open('test_out'+str(seed)+'.csv', 'w') as f: 95 | fileout = csv.writer(f) 96 | fileout.writerows(wavefuncs) 97 | with open('valid_out'+str(seed)+'.csv', 'w') as f: 98 | fileout = csv.writer(f) 99 | fileout.writerows(validfuncs) 100 | print 'Output complete' 101 | --------------------------------------------------------------------------------