Skip to content

Commit

Permalink
Add files via upload
Browse files Browse the repository at this point in the history
  • Loading branch information
@DurhamDecLab
DurhamDecLab committed Mar 3, 2019
1 parent 7b7204c commit a01312f
Showing 1 changed file with 49 additions and 37 deletions.
86 changes: 49 additions & 37 deletions ARBTools/ARBTrajec.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,22 +2,21 @@
import numpy as np
import random as rand
from ARBTools.ARBInterp import tricubic
import sys

#######################release the schmoo########################################
# Version 1.4, A4 code base

class trajectories:
def __init__(self, field, mass, moment):
def __init__(self, field):#, mass, moment):
self.eps = 10*np.finfo(float).eps # Machine precision of floating point number
self.amu = 1.66e-27 # Atomic mass units
self.k = 1.38064852e-23 # Boltzmann constant
mu = 9.274009994e-24 #Bohr magneton
self.m = mass*self.amu # Mass of atom
self.muom = moment*mu/self.m # magnetic moment in magnetons over mass
self.Interp = tricubic(field, mode='norm')
self.mu = 9.274009994e-24 #Bohr magneton
self.Interp = tricubic(field, 'quiet', mode='norm')

def atoms(self, T, N):
sd = np.sqrt(self.k*T/self.m)
def atoms(self, T, N, mass):
sd = np.sqrt(self.k*T/(mass*self.amu))
PSRange = np.zeros((int(N), 6))
for q in range(PSRange.shape[0]):
x = rand.gauss(0, 1e-3)
Expand All @@ -29,9 +28,11 @@ def atoms(self, T, N):
PSRange[q] = np.array([[x,y,z,vx,vy,vz]])
return PSRange

def sRK3d(self, sample, h): ### sample is the loaded distribution, h the RK timestep ###
def sRK3d(self, sample, h, mass, moment): ### sample is the loaded distribution, h the RK timestep ###
temp = sample.copy()

muom = moment*self.mu/mass # magnetic moment in magnetons over mass

### Runge-Kutta coefficients are stored in these arrays
RK0 = np.zeros((6))
RK1 = np.zeros((6))
Expand All @@ -42,11 +43,11 @@ def sRK3d(self, sample, h): ### sample is the loaded distribution, h t
Norm, Grad = self.Interp.sQuery(temp)

RK0[0] = h*(temp[3])
RK0[3] = -h*self.muom*Grad[0]
RK0[3] = -h*muom*Grad[0]
RK0[1] = h*(temp[4])
RK0[4] = -h*self.muom*Grad[1]
RK0[4] = -h*muom*Grad[1]
RK0[2] = h*(temp[5])
RK0[5] = -h*self.muom*Grad[2]
RK0[5] = -h*muom*Grad[2]

#New variables
temp = sample + 0.5*RK0
Expand All @@ -55,11 +56,11 @@ def sRK3d(self, sample, h): ### sample is the loaded distribution, h t
Norm, Grad = self.Interp.sQuery(temp)

RK1[0] = h*(temp[3])
RK1[3] = -h*self.muom*Grad[0]
RK1[3] = -h*muom*Grad[0]
RK1[1] = h*(temp[4])
RK1[4] = -h*self.muom*Grad[1]
RK1[4] = -h*muom*Grad[1]
RK1[2] = h*(temp[5])
RK1[5] = -h*self.muom*Grad[2]
RK1[5] = -h*muom*Grad[2]

#New variables
temp = sample + 0.5*RK1
Expand All @@ -68,11 +69,11 @@ def sRK3d(self, sample, h): ### sample is the loaded distribution, h t
Norm, Grad = self.Interp.sQuery(temp)

RK2[0] = h*(temp[3])
RK2[3] = -h*self.muom*Grad[0]
RK2[3] = -h*muom*Grad[0]
RK2[1] = h*(temp[4])
RK2[4] = -h*self.muom*Grad[1]
RK2[4] = -h*muom*Grad[1]
RK2[2] = h*(temp[5])
RK2[5] = -h*self.muom*Grad[2]
RK2[5] = -h*muom*Grad[2]

#New variables
temp = sample + RK2
Expand All @@ -81,19 +82,30 @@ def sRK3d(self, sample, h): ### sample is the loaded distribution, h t
Norm, Grad = self.Interp.sQuery(temp)

RK3[0] = h*(temp[3])
RK3[3] = -h*self.muom*Grad[0]
RK3[3] = -h*muom*Grad[0]
RK3[1] = h*(temp[4])
RK3[4] = -h*self.muom*Grad[1]
RK3[4] = -h*muom*Grad[1]
RK3[2] = h*(temp[5])
RK3[5] = -h*self.muom*Grad[2]
RK3[5] = -h*muom*Grad[2]

#Final new variables
temp = sample + (RK0 + 2*RK1 + 2*RK2 + RK3)/6
np.copyto(sample, temp)

def rRK3d(self, sample, h): ### sample is the loaded distribution, h the RK timestep ###
def rRK3d(self, sample, h, mass, moment): ### sample is the loaded distribution, h the RK timestep ###
temp = sample.copy()
N = int(temp.shape[0])
'''
if species=='Li':
mass = 6.941
moment = 1
elif species=='CaH':
mass = 41.085899
moment = 1
else:
sys.exit('--- Species not implemented ---')
'''
muom = moment*self.mu/mass # magnetic moment in magnetons over mass

### Runge-Kutta coefficients are stored in these arrays
RK0 = np.zeros((N,6))
Expand All @@ -105,11 +117,11 @@ def rRK3d(self, sample, h): ### sample is the loaded distribution, h t
Norms, Gradlist = self.Interp.rQuery(temp)

RK0[:, 0] = h*(temp[:, 3])
RK0[:, 3] = -h*self.muom*Gradlist[:, 0]
RK0[:, 3] = -h*muom*Gradlist[:, 0]
RK0[:, 1] = h*(temp[:, 4])
RK0[:, 4] = -h*self.muom*Gradlist[:, 1]
RK0[:, 4] = -h*muom*Gradlist[:, 1]
RK0[:, 2] = h*(temp[:, 5])
RK0[:, 5] = -h*self.muom*Gradlist[:, 2]
RK0[:, 5] = -h*muom*Gradlist[:, 2]

#New variables
temp = sample + 0.5*RK0
Expand All @@ -118,11 +130,11 @@ def rRK3d(self, sample, h): ### sample is the loaded distribution, h t
Norms, Gradlist = self.Interp.rQuery(temp)

RK1[:, 0] = h*(temp[:, 3])
RK1[:, 3] = -h*self.muom*Gradlist[:, 0]
RK1[:, 3] = -h*muom*Gradlist[:, 0]
RK1[:, 1] = h*(temp[:, 4])
RK1[:, 4] = -h*self.muom*Gradlist[:, 1]
RK1[:, 4] = -h*muom*Gradlist[:, 1]
RK1[:, 2] = h*(temp[:, 5])
RK1[:, 5] = -h*self.muom*Gradlist[:, 2]
RK1[:, 5] = -h*muom*Gradlist[:, 2]

#New variables
temp = sample + 0.5*RK1
Expand All @@ -131,11 +143,11 @@ def rRK3d(self, sample, h): ### sample is the loaded distribution, h t
Norms, Gradlist = self.Interp.rQuery(temp)

RK2[:, 0] = h*(temp[:, 3])
RK2[:, 3] = -h*self.muom*Gradlist[:, 0]
RK2[:, 3] = -h*muom*Gradlist[:, 0]
RK2[:, 1] = h*(temp[:, 4])
RK2[:, 4] = -h*self.muom*Gradlist[:, 1]
RK2[:, 4] = -h*muom*Gradlist[:, 1]
RK2[:, 2] = h*(temp[:, 5])
RK2[:, 5] = -h*self.muom*Gradlist[:, 2]
RK2[:, 5] = -h*muom*Gradlist[:, 2]

#New variables
temp = sample + RK2
Expand All @@ -144,32 +156,32 @@ def rRK3d(self, sample, h): ### sample is the loaded distribution, h t
Norms, Gradlist = self.Interp.rQuery(temp)

RK3[:, 0] = h*(temp[:, 3])
RK3[:, 3] = -h*self.muom*Gradlist[:, 0]
RK3[:, 3] = -h*muom*Gradlist[:, 0]
RK3[:, 1] = h*(temp[:, 4])
RK3[:, 4] = -h*self.muom*Gradlist[:, 1]
RK3[:, 4] = -h*muom*Gradlist[:, 1]
RK3[:, 2] = h*(temp[:, 5])
RK3[:, 5] = -h*self.muom*Gradlist[:, 2]
RK3[:, 5] = -h*muom*Gradlist[:, 2]

#Final new variables
temp = sample + (RK0 + 2*RK1 + 2*RK2 + RK3)/6
np.copyto(sample, temp)

def Iterate(self, sample, tm, h):
def Iterate(self, sample, tm, h, m, mom):
try:
if sample.shape[1] > 1:
t = 0
while t < tm:
self.rRK3d(sample, h) # Perform RK on distribution
self.rRK3d(sample, h, m, mom) # Perform RK on distribution
t += h
else:
t = 0
while t < tm:
self.sRK3d(sample, h) # Perform RK on distribution
self.sRK3d(sample, h, m, mom) # Perform RK on distribution
t += h
except IndexError:
t = 0
while t < tm:
self.sRK3d(sample, h) # Perform RK on distribution
self.sRK3d(sample, h, m, mom) # Perform RK on distribution
t += h

############################### Skookum ###############################

0 comments on commit a01312f

Please sign in to comment.