Homec4science
Diffusion PG2 export script (main) Browse

export_pg2/export_pg2.py
0382610b860fmain

Last Change
Raw File

export_pg2.py
View Options

#!/usr/bin/python3
# Note: change execution path above to use on other machines
# Copyright (c) 2021 Stephen Wasilewski, Lars O. Grobe HSLU
# =======================================================================
# This program is free software: you can redistribute it and/or
# modify it under the terms of theGNU Lesser General Public License
# as published by the Free Software Foundation, either version 3 of
# the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
# =======================================================================
import os
import shutil
import sys
import shlex
from clasp import click
import clasp.click_ext as clk
import clasp.script_tools as cst
import numpy as np
def row_c(row):
"""make tuple from table row data for checking against reference beam
Parameters
----------
row: list
[msid, sample label, mid, isRB, theta, phi, rho, tau, source,
source_filter, detector_filter]
Returns
-------
msid, source, source_filter, detector_filter
"""
return int(row[0]), row[8], row[9], row[10]
def get_rbeams(ttxt):
"""make reference beam dictionary and return measurement list"""
rbeams = {}
meas = []
for row in ttxt:
if row[3].lower() == "true":
rbeams[int(row[2])] = row_c(row)
else:
meas.append(row)
return rbeams, meas
def get_measurement_table(msids, rbid):
"""
Returns
-------
mtable: dict of lists
keys: sample labels
values: [mid, rbid, theta, phi]
vis: bool
is visible (based on source_filter == IRcutoff)
"""
pg_com = 'pgII_pstr --msid {} "%S %L %m %b %t %p %R %T %l %F %f" --ib'
p_mtable = []
rbeams = {}
print("* querying database for measurement table", file=sys.stderr)
for msid in msids:
ttxt = cst.pipeline([pg_com.format(msid)])
ttxt = [i.split() for i in ttxt.strip().split("\n")]
rbs, meas = get_rbeams(ttxt)
rbeams.update(rbs)
p_mtable += meas
mtable = {}
vis = {}
tp = {}
# reverses sort to prioritize newest, this means ultimately later samples
# are exported first (even though mids run in order within a sample)
for row in sorted(p_mtable, key=lambda x: -int(x[2])) :
if rbid in ['new', 'old', 'near']:
rbc = [k for k, v in rbeams.items() if v == row_c(row)]
if len(rbc) == 0:
raise ValueError("No valid reference beam for mid {}".format(row[2]))
if rbid == 'new':
rb = max(rbc)
elif rbid == 'old':
rb = min(rbc)
else:
rb = rbc[np.argmin(np.abs(np.array(rbc)-int(row[2])))]
else:
rb = rbid
if row[1] not in mtable:
mtable[row[1]] = []
vis[row[1]] = row[9] == "IRcutoff"
tp[row[1]] = []
tpk = f"{row[4]}_{row[5]}"
tpkr = len([i for i in tp[row[1]] if i == tpk])
if int(round(float(row[4]))) in [0, 180]:
tpkr = -tpkr - 1
line = (int(row[2]), rb, float(row[4]), float(row[5]), tpkr)
tp[row[1]] = [tpk] + tp[row[1]]
mtable[row[1]] = [line] + mtable[row[1]]
for sample in mtable.keys():
mtable[sample] = sorted(mtable[sample], key=lambda x: (x[2], x[3]))
return mtable, vis
def rbid_process(s):
"""parse rbid option"""
if s.lower() in ['new', 'old', 'near']:
return s.lower()
else:
try:
return int(s)
except ValueError:
raise ValueError("rbid must be an integer, 'new', 'old' or 'near'")
def get_flag(theta, opts):
"""output naming based on t/r and theta_in"""
if theta < 90 and "-t" in opts.split():
fl = "tfb"
elif "-t" in opts.split():
fl = "tbf"
elif theta < 90:
fl = "rff"
else:
fl = "rbb"
return fl
def plot_measurement(row, label="BIMSOLXXX", opts="-t", suffix=0,
tmpdir=".", overwrite=True):
"""
Parameters:
----------
row: list
[mid, rbid, theta, phi]
label: str
sample label
opts: str
opts for mountain, should include outfit filename
suffix: int
for sequential output naming / avoid clobbering
tmpdir: str
directory to write to
overwrite: bool
replace existing file
"""
fl = get_flag(row[2], opts)
suff = f"_p{suffix:03d}"
fname = f"{label}_{row[0]}_{row[2]:05.1f}_{row[3]:05.1f}_{fl}_DSF{suff}.jpg"
dcom = f"mountain -rb {row[1]} {row[0]} -f -e {opts} -geo 1024x768 -screendump {fname}"
if overwrite or not os.path.isfile(fname):
print(f"* running: {dcom}", file=sys.stderr)
cst.pipeline([dcom])
else:
print(f"* skipping: {fname} exists", file=sys.stderr)
return fname
def export_profile(row, label, fl, tmpdir=".", overwrite=True):
fname = f"{label}_{row[0]}_{row[2]:05.1f}_{row[3]:05.1f}_profile_{fl}_DSF.dat"
if fl[0] == 'r':
profile = "-r -ip 1,0,-90,90,0.5"
else:
profile = "-t -ip 1,0,90,270,0.5"
pcom = f"mountain -rb {row[1]} {row[0]} -f -e -n {profile},{fname}"
if overwrite or not os.path.isfile(fname):
print(f"* running: {pcom}", file=sys.stderr)
cst.pipeline([pcom])
else:
print(f"* skipping: {fname} exists", file=sys.stderr)
return fname
def deg(x):
return x*np.pi/180
def tp2xyz(th, ph):
x = np.sin(deg(th))*np.cos(deg(ph))
y = np.sin(deg(th))*np.sin(deg(ph))
z = np.cos(deg(th))
return x, y, z
def calc_dd_tau(th, ph, fname, tol=4):
header = ""
f = open(fname)
while True:
line = f.readline()
if line[0] == "#":
header += line
else:
break
dsf = np.loadtxt(fname)
xi, yi, zi = tp2xyz(th, ph)
xo, yo, zo = tp2xyz(dsf[:, 0], dsf[:, 1])
cosa = xi*xo*-1+yi*yo*-1+zi*zo*-1
alpha = np.arccos(cosa)*180/np.pi
dsf[alpha>tol/2, 2] = 0
tmp = fname.replace("DSF", "DSFDD")
np.savetxt(tmp, dsf, header=header, comments="")
result = shlex.split(cst.pipeline([f"mountain {tmp} -t -n"]))
os.remove(tmp)
return float(result[12])
def calc_dd_rho(th, ph, fname, tol=4):
header = ""
f = open(fname)
while True:
line = f.readline()
if line[0] == "#":
header += line
else:
break
dsf = np.loadtxt(fname)
xi, yi, zi = tp2xyz(th, ph)
xo, yo, zo = tp2xyz(dsf[:, 0], dsf[:, 1])
cosa = xi*xo*-1+yi*yo*-1+zi*zo
alpha = np.arccos(cosa)*180/np.pi
dsf[alpha>tol/2, 2] = 0
tmp = fname.replace("DSF", "DSFDD")
np.savetxt(tmp, dsf, header=header, comments="")
result = shlex.split(cst.pipeline([f"mountain {tmp} -t -n"]))
os.remove(tmp)
return float(result[12])
def export_measurement(row, label="BIMSOLXXX", opts=("-t", "-r"),
tmpdir=".", overwrite=True, plotviews=None,
dsf=True, profiles=True, shadesuff="orig",
exportall=False, direct=False):
"""
Parameters
----------
row: list
[mid, rbid, theta, phi]
label: str
sample label
opts: list
sets of options some combo of ['-t', '-r']
tmpdir: str
directory to write to
overwrite: bool
replace existing file
plotviews: list
oufit files (empty/None to skip plotting)
dsf: bool
datadump measurement to file
profiles: bool
interpolate profile data
shadesuff: str
suffix for self shaded reflection measurements
exportall: bool
include old measurements (replaced by refinement), suffix with odatXX
direct: bool
calculate direct-direct transmission and reflection for a 4 degree cone
"""
if plotviews is None:
plotviews = []
d = []
pr = []
pl = []
dd = []
if row[-1] not in [-1, 0] and not exportall:
return d, pr, pl, None
for opt in opts:
fl = get_flag(row[2], opt)
if fl[0] == 'r' and row[4] < 0:
suff = shadesuff
else:
suff = "dat"
if row[4] < -1:
suff = f"{suff}{abs(row[4])-1:02d}"
elif row[4] > 0:
suff = f"o{suff}{row[4]:02d}"
if dsf:
prefix = f"{tmpdir}/{label}/dsf/{fl}/{label}"
fname = f"{prefix}_{row[0]}_{row[2]:05.1f}_{row[3]:05.1f}_{fl}_DSF.{suff}"
dcom = f"mountain -rb {row[1]} {row[0]} -f -e -n -w {opt} -datadump {fname}"
d.append(fname)
if overwrite or not os.path.isfile(fname):
print(f"* running: {dcom}", file=sys.stderr)
cst.pipeline([dcom])
else:
print(f"* skipping: {fname} exists", file=sys.stderr)
if direct and fl[0] == 't':
dd.append(calc_dd_tau(row[2], row[3], fname))
elif direct:
dd.append(calc_dd_rho(row[2], row[3], fname))
# only export most recent (overwritten anyways in the naming)
if profiles and suff in ["dat", shadesuff]:
prefix = f"{tmpdir}/{label}/profiles/{fl}/{label}"
pr.append(export_profile(row, prefix, fl, tmpdir, overwrite))
for i, pv in enumerate(plotviews):
if suff in ["dat", shadesuff]:
prefix = f"{tmpdir}/{label}/plots/{fl}/{label}"
pl.append(plot_measurement(row, prefix, f"-{fl[0]} -fo {pv}", i,
tmpdir, overwrite))
if len(dd) == 0:
dd = None
return d, pr, pl, dd
def export_table(table, **kwargs):
"""calls export_measurement for a sample table"""
result = cst.pool_call(export_measurement, table, kwargs=kwargs)
return zip(*result)
def mkdirs(base, transmission, reflection):
"""setup output directories"""
cst.try_mkdir(base)
if transmission:
cst.try_mkdir(f"{base}/tfb")
cst.try_mkdir(f"{base}/tbf")
if reflection:
cst.try_mkdir(f"{base}/rff")
cst.try_mkdir(f"{base}/rbb")
def compress_folder(d, out, base):
print(f"* compressing: {d} to {out}", file=sys.stderr)
cst.pipeline([f"tar cjf {out} -C {base} {d}"])
def integral_table(table, label, transmission=True,
reflection=True, direct=None, vis=True, exportall=False):
if direct is None:
direct = [None] * len(table)
print(f"* writing table: {label}.txt", file=sys.stderr)
f = open(f"{label}.txt", "w")
f.write("# theta_in phi_in")
fstr = "%t %p"
if vis:
suf = "_vis"
else:
suf = ""
if transmission:
f.write(f" tau{suf}")
fstr += " %T"
if reflection:
f.write(f" rho{suf}")
fstr += " %R"
if direct and transmission:
f.write(f" tau_dir{suf}")
if direct and reflection:
f.write(f" rho_dir{suf}")
if exportall:
f.write(" flag (0:good, 1+:old, -1-:shaded reflection)")
else:
f.write(" flag (0:good, -1:shaded reflection)")
f.write("\n")
for row, dd in zip(table, direct):
if exportall or row[4] in [0, -1]:
ttxt = cst.pipeline([f"pgII_pstr --mid {row[0]} '{fstr}'"]).strip()
if dd is not None and transmission:
ttxt += f" {dd[0]}"
if dd is not None and reflection:
ttxt += f" {dd[-1]}"
fl = row[4]
f.write(f"{ttxt} {fl}\n")
f.close()
@click.command()
@click.argument('msid', callback=clk.split_int)
@click.option('-count', type=int, help="expected number of measurements. If "
"given will check for expected number of measurements. (not implemented)")
@click.option('-subset', callback=clk.split_int, help="for testing, give index of measurements to export")
@click.option('-shadesuff', default="orig",
help="suffix for data where arm blocks (typical suffix is dat)")
@click.option("--exportall/--no-exportall", default=False,
help="when True (default), exports only newest version of a "
"measurement. Otherwise exports old with a odatX suffix")
@click.option("-plotviews", default="", callback=clk.are_files,
help="if given, create mountain plots for given view"
" (the -fo option of mountain). If None, does not generate plots,"
" only data and sections (does not require x11)")
@click.option("-rbid", default='old', help="reference beam ID, (either a "
"number, 'new', 'old', or 'near' to automatically use newest, "
"oldest, or nearest within a session)")
@click.option("--tableonly/--no-tableonly", default=False,
help="skip export and only generate integral table (assumes they "
"are already calculated)")
@click.option("--transmission/--no-transmission", default=True,
help="export transmission")
@click.option("--reflection/--no-reflection", default=True,
help="export reflection")
@click.option("--dsf/--no-dsf", default=True,
help="export dsf data (warning, will not calculate integrals!)")
@click.option("--profiles/--no-profiles", default=True,
help="export interpolated profiles along phi 0")
@click.option("--overwrite/--no-overwrite", default=True,
help="if false, skip export of existing files")
@click.option("--direct/--no-direct", default=False,
help="calculate dir/dir rho/tau (not implemented)")
@click.option("-rotation", default=0.0,
help="rotation to apply (in degrees) to correct incident phi (not implemented)")
@click.option("-tmpdir", default="/var/scratch",
help="compressed output is saved to CWD, but export files "
"are written here")
@click.option("--deletetmp/--no-deletetmp", default=False,
help="delete output files after compression")
@click.option("--compress/--no-compress", default=True,
help="if False, overrides deletetmp")
@clk.shared_decs(clk.command_decs('0.1', wrap=True))
def main(ctx, msid, count=None, subset=None, shadesuff='orig', exportall=False, plotviews="",
rbid="old", transmission=True, reflection=True, dsf=True, tableonly=False, profiles=True, overwrite=True,
direct=False, rotation=0.0, tmpdir="/var/scratch", deletetmp=False, compress=True, **kwargs):
"""This script will export data and optionally plots for all measurement IDs
belonging to the given measurement session IDs. Consistency checks are
performed for the reference beam, shaded reflection measurements, and
optionally an expected measurement count.
For each sample label eg BIMSOLXXX, the following files are created::
- BIMSOLXXX.txt: table of: theta_in, phi_in, tau_hemi,
rho_hemi, tau_dir, rho_dir
- BIMSOLXXX_dsf.tar.bz2 containing for each measurement:
- BIMSOLXXX_<mid>_<theta>_<phi>_[tfb/tbf/rff/rbb]_DSF.dat
- BIMSOLXXX_plots.tar.bz2 containing for each measurement:
- BIMSOLXXX_<mid>_<theta>_<phi>_[tfb/tbf/rff/rbb]_DSF.jpg
- BIMSOLXXX_profiles.tar.bz2 containg for each measurement:
- BIMSOLXXX_<mid>_<theta>_<phi>_profile_[tfb/tbf/rff/rbb]_DSF.dat
Additionally, any hemispherical integrals calculated will be written back
to the database.
Arguments
---------
msid: list
the measurement session ids, give as space seperated (in quotes) list
of ints or range: "974 975 976" or "974:977" both yield:
[974, 975, 976]
"""
rbid = rbid_process(rbid)
mtable, vis = get_measurement_table(msid, rbid)
opts = []
if transmission:
opts.append("-t")
if reflection:
opts.append("-r")
for label, table in mtable.items():
if subset is not None:
table = [table[i] for i in subset]
cst.try_mkdir(f"{tmpdir}/{label}")
mkdirs(f"{tmpdir}/{label}/dsf", transmission, reflection)
if not tableonly:
if profiles:
mkdirs(f"{tmpdir}/{label}/profiles", transmission, reflection)
if len(plotviews) > 0:
mkdirs(f"{tmpdir}/{label}/plots", transmission, reflection)
d, pr, pl, dd = export_table(table, label=label, opts=opts,
plotviews=plotviews, tmpdir=tmpdir,
overwrite=overwrite, dsf=dsf,
profiles=profiles, shadesuff=shadesuff,
exportall=exportall, direct=direct)
if compress:
if dsf:
out = f"{label}_dsf.tar.bz2"
compress_folder(f"{label}/dsf", out, tmpdir)
if profiles:
out = f"{label}_profiles.tar.bz2"
compress_folder(f"{label}/profiles", out, tmpdir)
if len(plotviews) > 0:
out = f"{label}_plots.tar.bz2"
compress_folder(f"{label}/plots", out, tmpdir)
else:
dd = None
if direct and (tableonly or not dsf):
print("* Warning: direct values only computed when exporting"
" dsf data make sure --dsf is True and --tableonly is False"
" (--no-tableonly)", file=sys.stderr)
integral_table(table, label, transmission=transmission,
reflection=reflection, direct=dd, vis=vis[label],
exportall=exportall)
if deletetmp:
shutil.rmtree(f"{tmpdir}/{label}")
else:
print(f"clean up tmp files: rm -rf {tmpdir}/{label}", file=sys.stderr)
print("use --deletetmp to avoid this step", file=sys.stderr)
if __name__ == '__main__':
main()
c4science · Help