import os
import os.path
import numpy as np
import pickle
import astropy.units as u
import psycopg2
from solarsystem import planet_dist
import mathMB
from .satellite_initial_positions import satellite_initial_positions
from atomicdataMB import RadPresConst
from .LossInfo import LossInfo
from .rk5 import rk5
from .bouncepackets import bouncepackets
from .source_distribution import (surface_distribution, speed_distribution,
angular_distribution, surface_spot,
idlversion)
from mathMB import minmaxmean
[docs]class Output:
'''Keep track of packets initial and final positions
x0, y0, z0, f0, vx0, vy0, vz0
phi0, lat0, lon0
time, x, y, z, vx, vy, vz
index, npackets, totalsource
'''
def __init__(self, inputs, npackets, compress=True):
self.inputs = inputs
self.planet = inputs.geometry.planet
# Keep track of whether output is compressed
self.compress = compress
# Determine spatial unit
self.unit = u.def_unit('R_' + self.planet.object, self.planet.radius)
# Change unit for GM
self.GM = self.planet.GM.to(self.unit**3/u.s**2).value
# Determine distance and radial velocity of planet relative to the Sun
r, v_r = planet_dist(self.planet, self.inputs.geometry.taa)
self.aplanet = r.value
self.vrplanet = v_r.to(self.unit/u.s).value
# Find the default reactions and datasets
if inputs.options.lifetime.value <= 0:
self.loss_info = LossInfo(inputs.options.atom,
inputs.options.lifetime,
self.aplanet,
inputs.database)
else:
self.loss_info = None
# Set up the radiation pressure
if inputs.forces.radpres:
radpres = RadPresConst(inputs.options.atom, self.aplanet,
inputs.database)
radpres.velocity = radpres.velocity.to(self.unit/u.s).value
radpres.accel = radpres.accel.to(self.unit/u.s**2).value
self.radpres = radpres
else:
self.radpres = None
# Set up sticking
# if inputs.sticking_info['stickcoef'] != 1:
# stick = sticking_setup(inputs)
# Define the time that packets will run
if inputs.options.at_once:
time = np.ones(npackets) * inputs.options.endtime
else:
time = np.random.rand(npackets) * inputs.options.endtime
self.time0 = time.copy()
self.time = time.copy()
# Define the fractional content
self.f0 = np.ones(npackets)
self.frac = np.ones(npackets)
self.npackets = npackets
self.totalsource = np.sum(self.f0)
self.LossFrac = np.zeros(npackets)
# Determine initial satellite positions if necessary
if self.planet.moons is not None:
sat_init_pos = satellite_initial_positions(inputs)
else:
pass
# Determine initial source distribution
# Eventual way to do this
# if inputs.SpatialDist.type == 'inputfile':
# out = outputfile_source(inputs, packs_per_it, seed=seed)
# else:
# self.source_distribution()
self.source_distribution()
def __str__(self):
print('Contents of output:')
print('\tPlanet = {}'.format(self.planet.object))
print('\ta_planet = {}'.format(self.aplanet))
print('\tvr_planet = {}'.format(self.vrplanet))
# print(self.loss_info)
# print('\tNumber of Packets: {}'.format(self.npackets))
# print('\tUnits of time: {}'.format(self.time.unit))
# print('\tUnits of distance: {}'.format(self.X0.unit))
# print('\tUnits of velocity: {}'.format(self.V0.unit))
return ''
def __len__(self):
return self.npackets
def __getitem__(self, keys):
self.time = self.time[keys]
self.frac = self.frac[keys]
# self.LossFrac = self.LossFrac[keys]
self.x = self.x[keys]
self.y = self.y[keys]
self.z = self.z[keys]
self.vx = self.vx[keys]
self.vy = self.vy[keys]
self.vz = self.vz[keys]
if 'index' in self.__dict__.keys():
self.index = self.index[keys]
[docs] def source_distribution(self):
# Determine spatial distribution
if self.inputs.spatialdist.type == 'surface':
X0, lon, lat = surface_distribution(self.inputs,
self.npackets,
self.unit)
elif self.inputs.spatialdist.type == 'surfacespot':
X0, lon, lat = surface_spot(self.inputs,
self.npackets,
self.unit)
elif self.inputs.spatialdist.type == 'idlversion':
X0, V0_, _ = idlversion(self.inputs, self.unit)
else:
assert 0, 'Spatial Distribution {} not supported.'.format(
self.inputs.spatialdist.type)
# Choose a speed for each packet
v0 = speed_distribution(self.inputs, self.npackets)
if v0 is not None:
v0 = v0.to(self.unit/u.s)
# Choose direction for each packet
V0 = angular_distribution(self.inputs, X0.value, v0.value)
if V0 is not None:
V0 *= self.unit/u.s
else:
pass
else:
V0 = None
# Perturbation function
# Not installed yet
# Rotate everything to proper position for running the model
if (self.inputs.geometry.planet.object !=
self.inputs.geometry.startpoint):
assert 0, 'Not set up yet'
else:
pass
if V0 is None:
V0 = V0_
else:
pass
self.x0 = X0[0, :]
self.y0 = X0[1, :]
self.z0 = X0[2, :]
self.vx0 = V0[0, :]
self.vy0 = V0[1, :]
self.vz0 = V0[2, :]
self.x = X0[0, :]
self.y = X0[1, :]
self.z = X0[2, :]
self.vx = V0[0, :]
self.vy = V0[1, :]
self.vz = V0[2, :]
[docs] def driver(self):
# Set up the step sizes
hall = np.zeros(self.npackets) + 1000.
count = 0 # Number of steps taken
eps = 1e-10 # impact check delta
# These control how quickly the stepsize is increased or decreased
# between iterations
safety = 0.95
shrink = -0.25
grow = -0.2
# yscale = scaling parameter for each variable
# x,y,z ~ R_plan
# vx, vy, vz ~ 1 km/s (1/R_plan R_plan/s)
# frac ~ exp(-t/lifetime) ~ mean(frac)
rest = self.inputs.options.resolution
resx = self.inputs.options.resolution
resv = 0.1*self.inputs.options.resolution
resf = self.inputs.options.resolution
#########################################################
# Keep taking RK steps until every packet has reached the
# time of "image taken"
#########################################################
T = self.time.value
X = np.array([self.x, self.y, self.z])
V = np.array([self.vx, self.vy, self.vz])
Frac = self.frac
LossFrac = self.LossFrac
numb = []
moretogo = (T > rest) & (Frac > 0.)
while moretogo.any():
# Save old values
# This is used for determining if anything hits the rings
oldT = T[moretogo].copy()
oldX = X[:, moretogo].copy()
oldV = V[:, moretogo].copy()
oldFrac = Frac[moretogo].copy()
T0 = T[moretogo]
X0 = X[:, moretogo]
V0 = V[:, moretogo]
F0 = Frac[moretogo]
Loss0 = LossFrac[moretogo]
h = hall[moretogo]
assert np.all(h >= 0), '\n\tNegative values of h'
# Adjust stepsize to be no more than time remaining
h = (T0 >= h)*h + (T0 < h)*T0
# Run the rk5 step
t1, x1, v1, f1, delx, delv, delf = rk5(T0, X0, V0, F0, h, self)
# Do the error check
# scale = a_tol + |y|*r_tol
# for x: a_tol = r_tol = resolution
# for v: a_tol = r_tol = resolution/10.-require v more precise
# for f: a_tol = 0.01, r_tol = 0 -> frac tol = 1%
# -> f not implemented this way
scalex = resx + np.abs(x1)*resx
scalev = resv + np.abs(v1)*resv
scalef = resf + np.abs(f1)*resf
# Difference relative to acceptable difference
delx /= scalex
delv /= scalev
delf /= scalef
# Maximum error for each packet
xerrmax = delx.max(axis=0)
verrmax = delv.max(axis=0)
errmax = np.maximum(xerrmax, verrmax)
errmax = np.maximum(errmax, delf)
# error check
assert np.all(np.isfinite(errmax)), '\n\tInfinite values of emax'
# Make sure no negative frac
m = (f1 < 0) & (errmax < 1)
assert not np.any(m), 'Found new values of frac that are negative'
# Make sure frac doesn't increase
m = (f1-oldFrac > scalef) & (errmax > 1)
if np.any(m):
errmax[m] = 1.1
# Check where difference is very small. Adjust step size
noerr = errmax < 1e-7
errmax[noerr] = 1
h[noerr] *= 10
# Put the post-step values in
g = errmax < 1.0
b = (errmax >= 1.0)
numb.append(np.sum(b))
if np.any(g):
t_ = t1[g]
x_ = x1[:, g]
v_ = v1[:, g]
f_ = f1[g]
h_ = safety*h[g]*errmax[g]**grow
# Impact Check
tempR = np.linalg.norm(x_, axis=0)
# satrad = np.ones_like(tempR)
hhh = (tempR.value - 1) < 0
if np.any(hhh):
if self.inputs.sticking_info.stickcoef == 1.:
f_[hhh] = 0.
else:
bouncepackets(self, t_, x_, v_, f_, hhh)
# Check for escape
if not self.inputs.options.fullsystem:
f_[tempR > self.inputs.options.outeredge] = 0
# Check for vanishing
f_[f_ < 1e-10] = 0
# set remaining time = 0 for packets that are done
t_[f_ == 0] = 0.
# Put new values into arrays
T0[g] = t_
X0[:, g] = x_
V0[:, g] = v_
F0[g] = f_
h[g] = h_
if np.any(b):
# Don't adjust the bad value, but do fix the stepsize
htemp = safety*h[b]*errmax[b]**shrink
assert np.all(np.isfinite(htemp)), '\n\tInfinite values of h'
# Don't let step size drop below 1/10th previous step size
h[b] = np.maximum(htemp, 0.1*h[b])
assert np.all(h >= 0), '\n\tNegative values of h'
# Insert back into the original arrays
T[moretogo] = T0
X[:, moretogo] = X0
V[:, moretogo] = V0
Frac[moretogo] = F0
LossFrac[moretogo] = Loss0
hall[moretogo] = h
# Find which packets still need to run
moretogo = (T > rest) & (Frac > 0.)
if count % 100 == 0:
print(f'Step {count}. {np.sum(moretogo)} more to go\n'
f'\th: {mathMB.minmaxmean(h)}')
count += 1
# Put everything back into output
self.time = T*u.s
self.x = X[0, :]*self.unit
self.y = X[1, :]*self.unit
self.z = X[2, :]*self.unit
self.vx = V[0, :]*self.unit/u.s
self.vy = V[1, :]*self.unit/u.s
self.vz = V[2, :]*self.unit/u.s
self.frac = Frac
# self.LossFrac = lf0
# Add units back in
self.aplanet *= self.unit
self.vrplanet *= self.unit/u.s
self.GM *= self.unit**3/u.s**2
[docs] def stream_driver(self):
T = self.time.value
X = np.array([self.x, self.y, self.z])
V = np.array([self.vx, self.vy, self.vz])
Frac = self.frac
LossFrac = self.LossFrac
npack = len(T)
# Arrays to store the outputs
xx0 = np.zeros((3, npack, self.inputs.options.nsteps))
vv0 = np.zeros((3, npack, self.inputs.options.nsteps))
ff0 = np.zeros((npack, self.inputs.options.nsteps))
tt0 = np.zeros((npack, self.inputs.options.nsteps))
# Add initial values to the arrays
tt0[:, 0] = T
xx0[:, :, 0] = X
vv0[:, :, 0] = V
ff0[:, 0] = Frac
# step size and counters
dt = self.inputs.options.endtime.value/(self.inputs.options.nsteps-1)
h = dt
curtime = self.inputs.options.endtime.value
ct = 1
eps = 1e-10
moretogo = Frac > 0
while (curtime > 0) and (moretogo.any()):
t0 = T[moretogo]
x0 = X[:, moretogo]
v0 = V[:, moretogo]
f0 = Frac[moretogo]
oldF = f0.copy()
# Error checks
assert (f0 != 0).all(), 'Should not have 0 frac'
assert np.isfinite(x0).all(), 'Infinite X values'
assert np.isfinite(v0).all(), 'Infinite V values'
# Run the rk5 step
t1, x1, v1, f1, _, _, _ = rk5(t0, x0, v0, f0, h, self)
# Check for surface impacts
tempR = np.linalg.norm(x1, axis=0)
# satrad = np.ones_like(tempR)
hhh = (tempR - 1.) < 0
if np.any(hhh):
if self.inputs.sticking_info.stickcoef == 1:
f1[hhh] = 0.
else:
bouncepackets(self, t1, x1, v1, f1, hhh)
# Check for escape
if not self.inputs.options.fullsystem:
f1[tempR > self.inputs.options.outeredge] = 0
# Check to see if any packets have shrunk out of existence
f1[f1 < 1e-10] = 0.
# Set remaining time = 0 for packets that are done
t1[f1 == 0] = 0.
# Put new values back into the original array
T[moretogo] = t1
X[:, moretogo] = x1
V[:, moretogo] = v1
Frac[moretogo] = f1
# Save the results for later
tt0[:, ct] = T
xx0[:, :, ct] = X
vv0[:, :, ct] = V
ff0[:, ct] = Frac
# Check to see what still needs to be done
moretogo = Frac > 0
done = moretogo.any() is False
if (ct % 100) == 0:
print(ct, curtime, np.sum(moretogo))
# Update the times
ct += 1
curtime -= dt
# Put everything back into output
allpacks = npack*self.inputs.options.nsteps
self.time = tt0.reshape(allpacks)*u.s
self.x = xx0[0, :, :].reshape(allpacks)*self.unit
self.y = xx0[1, :, :].reshape(allpacks)*self.unit
self.z = xx0[2, :, :].reshape(allpacks)*self.unit
self.vx = vv0[0, :, :].reshape(allpacks)*self.unit/u.s
self.vy = vv0[1, :, :].reshape(allpacks)*self.unit/u.s
self.vz = vv0[2, :, :].reshape(allpacks)*self.unit/u.s
self.frac = ff0.reshape(allpacks)
# self.LossFrac = lf0
self.totalsource *= self.inputs.options.nsteps
index = np.mgrid[0:npack, 0:self.inputs.options.nsteps]
self.index = index[0, :, :].reshape(allpacks)
# Add units back in
self.aplanet *= self.unit
self.vrplanet *= self.unit/u.s
self.GM *= self.unit**3/u.s**2
[docs] def determine_filename(self):
'''Come up with a filename for the model'''
# TAA for observation
if self.planet.object == 'Mercury':
taastr = '{:03.0f}'.format(
np.round(self.inputs.geometry.taa.to(u.deg).value))
else:
assert 0, 'Filename not set up for anything but Mercury'
# Come up with a path name
pathname = os.path.join(self.inputs.savepath,
self.planet.object,
self.inputs.options.atom,
self.inputs.spatialdist.type,
self.inputs.speeddist.type,
taastr)
# Make the path if necessary
if os.path.exists(pathname) == 0:
os.makedirs(pathname)
numstr = '{:010d}'.format(self.idnum)
filename = f'{numstr}.pkl'
return pathname, filename
[docs] def save(self):
# Add output into database
con = psycopg2.connect(database=self.inputs.database)
con.autocommit = True
cur = con.cursor()
# Save filename and determine idnum
cur.execute('''INSERT into outputfile (npackets,
totalsource,
creationtime)
VALUES (%s, %s, NOW())''', (self.npackets,
self.totalsource))
cur.execute('''SELECT idnum
FROM outputfile
WHERE filename is Null''')
assert cur.rowcount == 1, 'Outputs with NULL filename'
self.idnum = cur.fetchone()[0]
# save the geometry
geometry = self.inputs.geometry
objs = [obj.object for obj in geometry.objects]
objs.sort()
phi = [p.value for p in geometry.phi]
cur.execute('''INSERT into geometry
VALUES (%s, %s, %s, %s::SSObject[],
%s, %s, '(%s, %s)', %s)''',
(self.idnum,
geometry.planet.object,
geometry.startpoint,
objs,
geometry.time,
phi,
geometry.subsolarpoint[0].value,
geometry.subsolarpoint[1].value,
geometry.taa.value))
con.commit()
# Save the sticking_info
sticking_info = self.inputs.sticking_info
cur.execute('''INSERT into sticking_info
VALUES (%s, %s, %s, %s, %s, %s,
%s, %s, %s, %s, %s)''',
(self.idnum,
sticking_info.stickcoef,
sticking_info.tsurf,
sticking_info.stickfn,
sticking_info.stick_mapfile,
sticking_info.epsilon,
sticking_info.n,
sticking_info.tmin,
sticking_info.emitfn,
sticking_info.accom_mapfile,
sticking_info.accom_factor))
con.commit()
# Save the forces
forces = self.inputs.forces
cur.execute('''INSERT into forces
VALUES (%s, %s, %s)''',
(self.idnum, forces.gravity, forces.radpres))
con.commit()
# Save the spatial distribution
spatdist = self.inputs.spatialdist
if spatdist.longitude is None:
lon = 0., 0.
else:
lon = spatdist.longitude[0].value, spatdist.longitude[1].value
if spatdist.latitude is None:
lat = 0., 0.
else:
lat = spatdist.latitude[0].value, spatdist.latitude[1].value
cur.execute('''INSERT into spatialdist
VALUES (%s, %s, %s, %s, %s,
ARRAY[%s,%s], ARRAY[%s,%s])''',
(self.idnum,
spatdist.type,
spatdist.exobase,
spatdist.use_map,
spatdist.mapfile,
lon[0], lon[1],
lat[0], lat[1]))
con.commit()
# Save the speed distribution
speeddist = self.inputs.speeddist
vprob = None if speeddist.vprob is None else speeddist.vprob.value
sigma = None if speeddist.sigma is None else speeddist.sigma.value
U = None if speeddist.U is None else speeddist.U.value
temperature = (None if speeddist.temperature is None
else speeddist.temperature.value)
delv = None if speeddist.delv is None else speeddist.delv.value
cur.execute('''INSERT into speeddist
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)''',
(self.idnum,
speeddist.type,
vprob,
sigma,
U,
speeddist.alpha,
speeddist.beta,
temperature,
delv))
con.commit()
# save the angular distribution
angdist = self.inputs.angulardist
if angdist.azimuth is None:
az = 'NULL'
else:
az0 = angdist.azimuth[0].value
az1 = angdist.azimuth[1].value
az = f'ARRAY[{az0}, {az1}]'
if angdist.altitude is None:
alt = 'NULL'
else:
alt0 = angdist.altitude[0].value
alt1 = angdist.altitude[1].value
alt = f'ARRAY[{alt0}, {alt1}]'
n = None if angdist.n is None else angdist.n
cur.execute(f'''INSERT into angulardist
VALUES (%s, %s, {az}, {alt}, %s)''',
(self.idnum,
angdist.type,
n))
con.commit()
# Save the options
options = self.inputs.options
cur.execute('''INSERT into options
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)''',
(self.idnum,
options.endtime.value,
options.resolution,
options.at_once,
options.atom,
options.lifetime.value,
options.fullsystem,
options.outeredge,
options.motion,
options.streamlines,
options.nsteps))
con.commit()
# Get a filename for the new savefile
# pathname, filename = self.determine_filename()
savefile = os.path.join(*self.determine_filename())
print('Saving output as {}'.format(savefile))
cur.execute('''UPDATE outputfile
SET filename=%s
WHERE idnum=%s''', (savefile, self.idnum))
con.commit()
con.close()
# Remove frac = 0
if self.compress:
self[self.frac > 0]
# Put the data into a better form for saving
self.time0 = self.time0.value.astype(np.float32)
self.f0 = self.f0.astype(np.float32)
self.x0 = self.x0.value.astype(np.float32)
self.y0 = self.y0.value.astype(np.float32)
self.z0 = self.z0.value.astype(np.float32)
self.vx0 = self.vx0.value.astype(np.float32)
self.vy0 = self.vy0.value.astype(np.float32)
self.vz0 = self.vz0.value.astype(np.float32)
self.time = self.time.value.astype(np.float32)
self.frac = self.frac.astype(np.float32)
self.x = self.x.value.astype(np.float32)
self.y = self.y.value.astype(np.float32)
self.z = self.z.value.astype(np.float32)
self.vx = self.vx.value.astype(np.float32)
self.vy = self.vy.value.astype(np.float32)
self.vz = self.vz.value.astype(np.float32)
self.LossFrac = self.LossFrac.astype(np.float32)
# Save output as a pickle
with open(savefile, 'wb') as f:
pickle.dump(self, f, protocol=pickle.HIGHEST_PROTOCOL)
[docs] @classmethod
def restore(cls, filename, reform=False):
output = pickle.load(open(filename, 'rb'))
output.time0 *= u.s
output.time *= u.s
output.x0 *= output.unit
output.y0 *= output.unit
output.z0 *= output.unit
output.vx0 *= output.unit/u.s
output.vy0 *= output.unit/u.s
output.vz0 *= output.unit/u.s
output.x *= output.unit
output.y *= output.unit
output.z *= output.unit
output.vx *= output.unit/u.s
output.vy *= output.unit/u.s
output.vz *= output.unit/u.s
if reform:
index_ = list(set(output.index))
n, m = len(index_), output.inputs.options.nsteps
x = np.ndarray((n, m))
y = np.ndarray((n, m))
z = np.ndarray((n, m))
vx = np.ndarray((n, m))
vy = np.ndarray((n, m))
vz = np.ndarray((n, m))
frac, index = np.ndarray((n, m)), np.ndarray((n, m))
time = np.ndarray((n, m))
for i, ind in enumerate(index_):
q = output.index == ind
x[i, :sum(q)] = output.x[q]
y[i, :sum(q)] = output.y[q]
z[i, :sum(q)] = output.z[q]
vx[i, :sum(q)] = output.vx[q]
vy[i, :sum(q)] = output.vy[q]
vz[i, :sum(q)] = output.vz[q]
frac[i, :sum(q)] = output.frac[q]
index[i, :sum(q)] = output.index[q]
time[i, :sum(q)] = output.time[q]
output.x = x
output.y = y
output.z = z
output.vx = vx
output.vy = vy
output.vz = vz
output.frac = frac
output.index = index
output.time = time
return output