diff --git a/tofu/__init__.py b/tofu/__init__.py index 4ef97c747..7bf166d8c 100644 --- a/tofu/__init__.py +++ b/tofu/__init__.py @@ -74,7 +74,7 @@ # ------------------------------------- msg = None -dsub = dict.fromkeys(['imas2tofu', 'mag']) +dsub = dict.fromkeys(['imas2tofu', 'openadas2tofu', 'mag']) for sub in dsub.keys(): try: exec('import tofu.{0} as {0}'.format(sub)) diff --git a/tofu/openadas2tofu/__init__.py b/tofu/openadas2tofu/__init__.py new file mode 100644 index 000000000..8daa3885c --- /dev/null +++ b/tofu/openadas2tofu/__init__.py @@ -0,0 +1,21 @@ +# -*- coding: utf-8 -*- +#! /usr/bin/python +""" +The openadas-compatibility module of tofu + +""" +import traceback + +try: + try: + from tofu.openadas2tofu._requests import * + from tofu.openadas2tofu._read_files import * + except Exception: + from ._requests import * + from ._read_files import * +except Exception as err: + msg = str(traceback.format_exc()) + msg += "\n\n => the optional sub-package tofu.imas2tofu is not usable\n" + raise Exception(msg) + +del traceback diff --git a/tofu/openadas2tofu/_read_files.py b/tofu/openadas2tofu/_read_files.py new file mode 100644 index 000000000..cebe3a623 --- /dev/null +++ b/tofu/openadas2tofu/_read_files.py @@ -0,0 +1,529 @@ + +# Built-in +import os +import re +import itertools as itt + +# Common +import numpy as np +from scipy.interpolate import RectBivariateSpline as scpRectSpl + + +__all__ = ['read', 'read_all'] + + +_DTYPES = {'adf11': ['acd', 'ccd', 'scd', 'plt', 'prb'], + 'adf15': None} +_DEG = 1 + + +# ############################################################################# +# Utility functions +# ############################################################################# + + +def _get_PATH_LOCAL(): + pfe = os.path.join(os.path.expanduser('~'), '.tofu', 'openadas2tofu') + if os.path.isdir(pfe): + return pfe + else: + return None + + +def _get_subdir_from_pattern(path, pattern): + ld = [dd for dd in os.listdir(path) + if (os.path.isdir(os.path.join(path, dd)) + and pattern in dd)] + if len(ld) != 1: + av = [dd for dd in os.listdir(path) + if os.path.isdir(os.path.join(path, dd))] + msg = ("You have no / many directories in your local " + + "~/.tofu/openadas2tofu/ matching the desired file type:\n" + + "\t- provided : {}\n".format(pattern) + + "\t- available: {}\n".format(av) + + " => download the data with tf.openadas2tofu.download()") + raise Exception(msg) + return os.path.join(path, ld[0]) + + +# ############################################################################# +# Main functions +# ############################################################################# + + +def read(adas_path, **kwdargs): + """ Read openadas-formatted files and return a dict with the data + + Povide the full adas file name + The result is returned as a dict + + example + ------- + >>> import tofu as tf + >>> fn = '/adf11/scd74/scd74_ar.dat' + >>> out = tf.openadas2tofu.read(fn) + >>> fn = '/adf15/pec40][ar/pec40][ar_ca][ar16.dat' + >>> out = tf.openadas2tofu.read(fn) + """ + + path_local = _get_PATH_LOCAL() + + # Check whether the local .tofu repo exists, if not recommend tofu-custom + if path_local is None: + path = os.path.join(os.path.expanduser('~'), '.tofu', 'openadas2tofu') + msg = ("You do not seem to have a local ./tofu repository\n" + + "tofu uses that local repository to store all user-specific " + + "data and downloads\n" + + "In particular, openadas files are downloaded and saved in:\n" + + "\t{}\n".format(path) + + " => to set-up your local .tofu repo, run in a terminal:\n" + + "\ttofu-custom") + raise Exception(msg) + + # Determine whether adas_path is an absolute path or an adas full name + if os.path.isfile(adas_path): + pfe = adas_path + else: + # make sure adas_path is not understood as absolute local path + if adas_path[0] == '/': + adas_path = adas_path[1:] + + # Check file was downloaded locally + pfe = os.path.join(path_local, adas_path) + if not os.path.isfile(pfe): + msg = ("Provided file does not seem to exist:\n" + + "\t{}\n".format(pfe) + + " => Search it online with tofu.openadas2tofu.search()\n" + + " => Download it with tofu.openadas2tofu.download()") + raise FileNotFoundError(msg) + + lc = [ss for ss in _DTYPES.keys() if ss in pfe] + if not len(lc) == 1: + msg = ("File type could not be derived from absolute path:\n" + + "\t- provided: {}\n".format(pfe) + + "\t- supported: {}".format(sorted(_DTYPES.keys()))) + raise Exception(msg) + + func = eval('_read_{}'.format(lc[0])) + return func(pfe, **kwdargs) + + +def read_all(element=None, charge=None, typ1=None, typ2=None, + verb=None, **kwdargs): + """ Read all relevant openadas files for chosen typ1 + + Please specify: + - typ1: 'adf11' or 'adf15' + - element: the symbol of the element + + If typ1 = 'adf11', you can also provide typ2 to specify the coefficients: + - 'scd': effective ionisation coefficients + - 'acd': effective electron-impact recombination coefficients + - 'ccd': effective hydrogen-impact recombination coefficients + - 'plt': line power due to electron-impact excitation + - 'prc': line power due to hydrogen-impact excitation + - 'prb': rad. recombination and bremmstrahlung due to electron-impact + + If typ1 = 'adf15', you can optioanlly provide a min/max wavelength + + The result is returned as a dict + + examples + -------- + >>> import tofu as tf + >>> dout = tf.openadas2tofu.read_all(element='ar', typ1='adf11') + >>> dout = tf.openadas2tofu.read_all(element='ar', typ1='adf15', + charge=16, + lambmin=3.94e-10, lambmax=4.e-10) + """ + + # -------------------- + # Check whether the local .tofu repo exists, if not recommend tofu-custom + path_local = _get_PATH_LOCAL() + if path_local is None: + path = os.path.join(os.path.expanduser('~'), '.tofu', 'openadas2tofu') + msg = ("You do not seem to have a local ./tofu repository\n" + + "tofu uses that local repository to store all user-specific " + + "data and downloads\n" + + "In particular, openadas files are downloaded and saved in:\n" + + "\t{}\n".format(path) + + " => to set-up your local .tofu repo, run in a terminal:\n" + + "\ttofu-custom") + raise Exception(msg) + + # -------------------- + # Check / format input + if typ1 is None: + typ1 = 'adf15' + if not isinstance(typ1, str) or typ1.lower() not in _DTYPES.keys(): + msg = ("Please choose a valid adas file type:\n" + + "\t- allowed: {}\n".format(_DTYPES.keys()) + + "\t- provided: {}".format(typ1)) + raise Exception(msg) + typ1 = typ1.lower() + if typ1 == 'adf11' and typ2 is None: + typ2 = _DTYPES[typ1] + fd = os.listdir(os.path.join(path_local, typ1)) + typ2 = [sorted([ss for ss in fd if tt in ss])[-1] for tt in typ2] + if isinstance(typ2, str): + typ2 = [typ2] + if (_DTYPES[typ1] is not None + and (not isinstance(typ2, list) + or not all([any([s1 in ss for s1 in _DTYPES[typ1]]) + for ss in typ2]))): + msg = ("typ2 must be a list of valid openadas file types for typ1:\n" + + "\t- provided: {}\n".format(typ2) + + "\t- available for {}: {}".format(typ1, _DTYPES[typ1])) + raise Exception(msg) + + if not isinstance(element, str): + msg = "Please choose an element!" + raise Exception(msg) + element = element.lower() + if charge is not None and not isinstance(charge, int): + msg = "charge must be a int!" + raise Exception(msg) + + if verb is None: + verb = True + + # -------------------- + # Get list of relevant directories + + # Level 1: Type + path = _get_subdir_from_pattern(path_local, typ1) + # Level 2: element or typ2 + if typ1 == 'adf11': + lpath = [_get_subdir_from_pattern(path, tt) for tt in typ2] + elif typ1 == 'adf15': + lpath = [_get_subdir_from_pattern(path, element)] + + # -------------------- + # Get list of relevant files pfe + lpfe = list(itt.chain.from_iterable( + [[os.path.join(path, ff) for ff in os.listdir(path) + if (os.path.isfile(os.path.join(path, ff)) + and ff[-4:] == '.dat' + and element in ff)] + for path in lpath])) + + if charge is not None and typ1 == 'adf15': + lpfe = [ff for ff in lpfe if str(charge) in ff] + + # -------------------- + # Extract data from each file + func = eval('_read_{}'.format(typ1)) + dout = {} + for pfe in lpfe: + if verb is True: + msg = "\tLoading data from {}".format(pfe) + print(msg) + out = func(pfe, dout=dout, **kwdargs) + return out + + +# ############################################################################# +# Specialized functions for ADF 11 +# ############################################################################# + + +def _read_adf11(pfe, deg=None, dout=None): + if deg is None: + deg = _DEG + if dout is None: + dout = {} + + # Get second order file type + typ1 = [vv for vv in _DTYPES['adf11'] if vv in pfe] + if len(typ1) != 1: + msg = ("Second order file type culd not be inferred from file name!\n" + + "\t- available: {}\n".format(_DTYPES['adf11']) + + "\t- provided: {}".format(pfe)) + raise Exception(msg) + typ1 = typ1[0] + + # Get element + elem = pfe[:-4].split('_')[1] + comline = '-'*60 + comline2 = 'C'+comline + + if typ1 in ['acd', 'ccd', 'scd', 'plt', 'prb']: + + # read blocks + with open(pfe) as search: + for ii, line in enumerate(search): + + if comline2 in line: + break + + # Get atomic number (transitions) stored in this file + if ii == 0: + lstr = line.split('/') + lin = [ss for ss in lstr[0].strip().split(' ') + if ss.strip() != ''] + lc = [len(lin) == 5 and all([ss.isdigit() for ss in lin]), + elem.upper() in lstr[1], + 'ADF11' in lstr[2]] + if not all(lc): + msg = ("File header format seems to have changed!\n" + + "\t- lc = {}".format(lc)) + raise Exception(msg) + Z, nne, nte, q0, qend = map(int, lin) + nelog10 = np.array([]) + telog10 = np.array([]) + in_ne = True + continue + + if comline in line: + continue + + # Get nelog10 + if in_ne: + li = [ss for ss in line.strip().split(' ') + if ss.strip() != ''] + nelog10 = np.append(nelog10, np.array(li, dtype=float)) + if nelog10.size == nne: + in_ne = False + in_te = True + + # Get telog10 + elif in_te is True: + li = [ss for ss in line.strip().split(' ') + if ss.strip() != ''] + telog10 = np.append(telog10, np.array(li, dtype=float)) + if telog10.size == nte: + in_te = False + in_ion = True + + # Get ion block + elif (in_ion is True and 'Z1=' in line + and ('------/' in line and 'DATE=' in line)): + nion = int( + line[line.index('Z1=')+len('Z1='):].split('/')[0]) + if typ1 in ['scd', 'plt']: + charge = nion - 1 + else: + charge = nion + coefslog10 = np.array([]) + elif in_ion is True and charge is not None: + li = [ss for ss in line.strip().split(' ') + if ss.strip() != ''] + coefslog10 = np.append(coefslog10, + np.array(li, dtype=float)) + if coefslog10.size == nne*nte: + key = '{}{}'.format(elem, charge) + tkv = [('element', elem), ('Z', Z), ('charge', charge)] + if key in dout.keys(): + assert all([dout[key][ss] == vv for ss, vv in tkv]) + else: + dout[key] = {ss: vv for ss, vv in tkv} + if typ1 == 'scd': + # nelog10+6 to convert /cm3 -> /m3 + # coefslog10-6 to convert cm3/s -> m3/s + func = scpRectSpl(nelog10+6, telog10, + coefslog10.reshape((nne, nte))-6, + kx=deg, ky=deg) + dout[key]['ionis'] = {'func': func, + 'type': 'log10_nete', + 'units': 'log10(m3/s)', + 'source': pfe} + elif typ1 == 'acd': + # nelog10+6 to convert /cm3 -> /m3 + # coefslog10-6 to convert cm3/s -> m3/s + func = scpRectSpl(nelog10+6, telog10, + coefslog10.reshape((nne, nte))-6, + kx=deg, ky=deg) + dout[key]['recomb'] = {'func': func, + 'type': 'log10_nete', + 'units': 'log10(m3/s)', + 'source': pfe} + elif typ1 == 'ccd': + # nelog10+6 to convert /cm3 -> /m3 + # coefslog10-6 to convert cm3/s -> m3/s + func = scpRectSpl(nelog10+6, telog10, + coefslog10.reshape((nne, nte))-6, + kx=deg, ky=deg) + dout[key]['recomb_ce'] = {'func': func, + 'type': 'log10_nete', + 'units': 'log10(m3/s)', + 'source': pfe} + elif typ1 == 'plt': + # nelog10+6 to convert /cm3 -> /m3 + # coefslog10+6 to convert W.cm3 -> W.m3 + func = scpRectSpl(nelog10+6, telog10, + coefslog10.reshape((nne, nte))+6, + kx=deg, ky=deg) + dout[key]['rad_bb'] = {'func': func, + 'type': 'log10_nete', + 'units': 'log10(W.m3)', + 'source': pfe} + elif typ1 == 'prb': + # nelog10+6 to convert /cm3 -> /m3 + # coefslog10+6 to convert W.cm3 -> W.m3 + func = scpRectSpl(nelog10+6, telog10, + coefslog10.reshape((nne, nte))+6, + kx=deg, ky=deg) + dout[key]['rad_fffb'] = {'func': func, + 'type': 'log10_nete', + 'units': 'log10(W.m3)', + 'source': pfe} + if nion == Z: + break + + return dout + + +# ############################################################################# +# Specialized functions for ADF 15 +# ############################################################################# + + +def _get_adf15_key(elem, charge, isoel, typ0, typ1): + return '{}{}_{}_openadas_{}_{}'.format(elem, charge, isoel, + typ0, typ1) + + +def _read_adf15(pfe, dout=None, + lambmin=None, + lambmax=None, + deg=None): + + if deg is None: + deg = _DEG + if dout is None: + dout = {} + + # Get summary of transitions + flagblock = '/isel =' + flag0 = 'superstage partition information' + + # Get file markers from name (elem, charge, typ0, typ1) + typ0, typ1, elemq = pfe.split('][')[1:] + ind = re.search(r'\d', elemq).start() + elem = elemq[:ind].title() + charge = int(elemq[ind:-4]) + assert elem.lower() in typ0[:2] + assert elem.lower() == typ1.split('_')[0] + typ0 = typ0[len(elem)+1:] + typ1 = typ1.split('_')[1] + + # Extract data from file + nlines, nblock = None, 0 + in_ne, in_te, in_pec, in_tab, itab = False, False, False, False, np.inf + skip = False + with open(pfe) as search: + for ii, line in enumerate(search): + + # Get number of lines (transitions) stored in this file + if ii == 0: + lstr = line.split('/') + nlines = int(lstr[0].replace(' ', '')) + continue + + # Get info about the transition being scanned (block) + if flagblock in line and 'C' not in line and nblock < nlines: + lstr = [kk for kk in line.rstrip().split(' ') if len(kk) > 0] + lamb = float(lstr[0])*1.e-10 + isoel = nblock + 1 + nblock += 1 + c0 = ((lambmin is not None and lamb < lambmin) + or (lambmax is not None and lamb > lambmax)) + if c0: + skip = True + continue + skip = False + nne, nte = int(lstr[1]), int(lstr[2]) + typ = [ss[ss.index('type=')+len('type='):ss.index('/ispb')] + for ss in lstr[3:] if 'type=' in ss] + assert len(typ) == 1 + # To be updated : proper rezading from line + in_ne = True + ne = np.array([]) + te = np.array([]) + pec = np.full((nne*nte,), np.nan) + ind = 0 + continue + + if 'root partition information' in line and skip is True: + skip = False + + # Check lamb is ok + if skip is True: + continue + + # Get ne for the transition being scanned (block) + if in_ne is True: + ne = np.append(ne, + np.array(line.rstrip().strip().split(' '), + dtype=float)) + if ne.size == nne: + in_ne = False + in_te = True + + # Get te for the transition being scanned (block) + elif in_te is True: + te = np.append(te, + np.array(line.rstrip().strip().split(' '), + dtype=float)) + if te.size == nte: + in_te = False + in_pec = True + + # Get pec for the transition being scanned (block) + elif in_pec is True: + data = np.array(line.rstrip().strip().split(' '), + dtype=float) + pec[ind:ind+data.size] = data + ind += data.size + if ind == pec.size: + in_pec = False + key = _get_adf15_key(elem, charge, isoel, typ0, typ1) + # log(ne)+6 to convert /cm3 -> /m3 + # log(pec)+6 to convert cm3/s -> m3/s + func = scpRectSpl(np.log(ne)+6, np.log(te), + np.log(pec).reshape((nne, nte))+6, + kx=deg, ky=deg) + dout[key] = {'lambda': lamb, + 'ION': '{}{}+'.format(elem, charge), + 'symbol': '{}{}-{}'.format(typ0, typ1, isoel), + 'origin': pfe, + 'type': typ[0], + 'ne': ne, 'te': te, + 'pec': {'func': func, + 'type': 'log_nete', + 'units': 'log(m3/s)', + 'source': 'pfe'}} + + # Get transitions from table at the end + if 'photon emissivity atomic transitions' in line: + itab = ii + 6 + if ii == itab: + in_tab = True + if in_tab is True: + lstr = [kk for kk in line.rstrip().split(' ') if len(kk) > 0] + isoel = int(lstr[1]) + lamb = float(lstr[2])*1.e-10 + key = _get_adf15_key(elem, charge, isoel, typ0, typ1) + c0 = ((lambmin is None or lambmin < lamb) + and (lambmax is None or lambmax > lamb)) + if c0 and key not in dout.keys(): + msg = ("Inconsistency in file {}:\n".format(pfe) + + "\t- line should be present".format(key)) + raise Exception(msg) + if key in dout.keys(): + if dout[key]['lambda'] != lamb: + msg = "Inconsistency in file {}".format(pfe) + raise Exception(msg) + c0 = (dout[key]['type'] not in lstr + or lstr.index(dout[key]['type']) < 4) + if c0: + msg = ("Inconsistency in table, type not found:\n" + + "\t- expected: {}\n".format(dout[key]['type']) + + "\t- line: {}".format(line)) + raise Exception(msg) + trans = lstr[3:lstr.index(dout[key]['type'])] + dout[key]['transition'] = ''.join(trans) + if isoel == nlines: + in_tab = False + assert all(['transition' in vv.keys() for vv in dout.values()]) + return dout diff --git a/tofu/openadas2tofu/_requests.py b/tofu/openadas2tofu/_requests.py new file mode 100644 index 000000000..ab1f24ebb --- /dev/null +++ b/tofu/openadas2tofu/_requests.py @@ -0,0 +1,486 @@ + +# Built-in +import os +import shutil +import requests +import warnings + +# Common +import numpy as np + + +__all__ = ['search_online', 'search_online_by_wavelengthA', + 'download', 'download_all', 'clean_downloads'] + + +# Check whether a local .tofu/ repo exists +_URL = 'https://open.adas.ac.uk' +_URL_SEARCH = _URL + '/freeform?searchstring=' +_URL_SEARCH_WAVL = _URL + '/wavelength?' +_URL_ADF15 = _URL + '/adf15' +_URL_DOWNLOAD = _URL + '/download' + +_INCLUDE_PARTIAL = True + + +# ############################################################################# +# Utility functions +# ############################################################################# + + +def _get_PATH_LOCAL(): + pfe = os.path.join(os.path.expanduser('~'), '.tofu', 'openadas2tofu') + if os.path.isdir(pfe): + return pfe + else: + return None + + +def _getcharray(ar, col=None, sep=' ', line='-', just='l', + verb=True, returnas=str): + """ Format and return char array (for pretty printing) """ + c0 = ar is None or len(ar) == 0 + if c0: + return '' + ar = np.array(ar, dtype='U') + + if ar.ndim == 1: + ar = ar.reshape((1, ar.size)) + + # Get just len + nn = np.char.str_len(ar).max(axis=0) + if col is not None: + if len(col) not in ar.shape: + msg = ("len(col) should be in np.array(ar, dtype='U').shape:\n" + + "\t- len(col) = {}\n".format(len(col)) + + "\t- ar.shape = {}".format(ar.shape)) + raise Exception(msg) + if len(col) != ar.shape[1]: + ar = ar.T + nn = np.char.str_len(ar).max(axis=0) + nn = np.fmax(nn, [len(cc) for cc in col]) + + # Apply to array + fjust = np.char.ljust if just == 'l' else np.char.rjust + out = np.array([sep.join(v) for v in fjust(ar, nn)]) + + # Apply to col + if col is not None: + arcol = np.array([col, [line*n for n in nn]], dtype='U') + arcol = np.array([sep.join(v) for v in fjust(arcol, nn)]) + out = np.append(arcol, out) + + if verb is True: + print('\n'.join(out)) + if returnas is str: + return '\n'.join(out) + elif returnas is np.ndarray: + return ar + + +# ############################################################################# +# online search +# ############################################################################# + + +class FileAlreayExistsException(Exception): + pass + + +def search_online(searchstr=None, returnas=None, + include_partial=None, verb=None): + """ Perform an online freeform search on https://open.adas.ac.uk + + Pass searchstr to the online freeform search + Prints the results (if verb=True) + Optionally return the result as: + - a char array (returnas = np.ndarray) + - a formatted str (returnas = str) + + example + ------- + >>> import tofu as tf + >>> tf.openadas2tofu.search_online('ar+16 ADF15') + """ + + # Check input + if returnas is None: + returnas = False + if verb is None: + verb = True + if include_partial is None: + include_partial = _INCLUDE_PARTIAL + if searchstr is None: + searchstr = '' + searchurl = '+'.join([requests.utils.quote(kk) + for kk in searchstr.split(' ')]) + + total_url = '{}{}&{}'.format(_URL_SEARCH, searchurl, 'searching=1') + resp = requests.get(total_url) + + # Extract response from html + out = resp.text.split('\n') + flag0 = '' + flag1 = '
' + ind0 = [ii for ii, vv in enumerate(out) if flag0 in vv] + ind1 = [ii for ii, vv in enumerate(out) if flag1 in vv] + if len(ind0) != 1 or len(ind1) == 0: + msg = ("Format of html response seems to have changed!\n" + + "Cannot find flags:\n" + + "\t- {}\n".format(flag0) + + "\t- {}\n".format(flag1) + + "in requests.get({}).text".format(total_url)) + raise Exception(msg) + ind1 = np.min([ii for ii in ind1 if ii > ind0[0]]) + out = out[ind0[0] + 1:ind1-1] + nresults = len(out) - 1 + + # Get columns + heads = [str.replace(kk.replace('', '').replace('', ''), + '', '').replace('', '') + for kk in out[0].split('')] + nhead = len(heads) + + # Get results + lout = [] + for ii in range(0, nresults): + if 'Partial results are listed below' in out[ii+1]: + if include_partial is False: + break + lout.append(['-', '-', '----- (partial results) -----', '-']) + ind = out[ii+1].index('below') + len('below') + out[ii+1] = out[ii+1][ind:] + lstri = out[ii+1].split('', '') + if '' in elmq: + elm = elmq[:elmq.index('')] + charge = elmq[elmq.index('')+len(''):] + charge = charge.replace('', '') + if charge == '+': + charge = '1+' + else: + charge = '' + elm = elmq + typ = lstri[2][1:] + fil = lstri[3][lstri[3].index('detail')+len('detail'):] + fil = fil[:fil.index('.dat')+len('.dat')] + lout.append([elm, charge, typ, fil]) + + # Format output + char = np.array(lout) + col = ['Element', 'charge', 'type of data', 'full file name'] + arr = _getcharray(char, col=col, + sep=' ', line='-', just='l', + returnas=returnas, verb=verb) + return arr + + +def search_online_by_wavelengthA(lambmin=None, lambmax=None, resolveby=None, + element=None, charge=None, + returnas=None, verb=None): + """ Perform an online search by wavelength on https://open.adas.ac.uk + + Pass the min / max wavelength (in Angstrom) to the online wavelength search + Prints the results (if verb=True) + Optionally return the result as: + - a char array (returnas = np.ndarray) + - a formatted str (returnas = str) + + The result can be resolve by transition or by adas file + Optionally filter by element to return only the results of one element + + example + ------- + >>> import tofu as tf + >>> tf.openadas2tofu.search_online_by_wavelengthA(3., 4., element='ar') + """ + + # Check input + if returnas is None: + returnas = False + if verb is None: + verb = True + if lambmin is None: + lambmin = '' + if lambmax is None: + lambmax = '' + if resolveby is None: + resolveby = 'transition' + if resolveby not in ['transition', 'file']: + msg = ("Arg resolveeby must be:\n" + + "\t- 'transition': list all available transitions\n" + + "\t- 'file': list all files containing relevant transitions") + raise Exception(msg) + if element is not None and not isinstance(element, str): + msg = ("Arg element must be a str (e.g.: element='ar')\n" + + "\t- provided: {}".format(element)) + raise Exception(msg) + if charge is not None: + if not isinstance(charge, int): + msg = ("Arg charge must be a int!\n" + + "\t- provided: {}".format(charge)) + raise Exception(msg) + charge = '0' if charge == 0 else '{}+'.format(charge) + + searchurl = '&'.join(['wave_min={}'.format(lambmin), + 'wave_max={}'.format(lambmax), + 'resolveby={}'.format(resolveby)]) + + total_url = '{}{}&{}'.format(_URL_SEARCH_WAVL, searchurl, 'searching=1') + resp = requests.get(total_url) + + # Extract response from html + out = resp.text.split('\n') + flag0 = '' + flag1 = '
' + ind0 = [ii for ii, vv in enumerate(out) if flag0 in vv] + ind1 = [ii for ii, vv in enumerate(out) if flag1 in vv] + if len(ind0) != 1 or len(ind1) != 1: + msg = ("Format of html response seems to have changed!\n" + + "Cannot find flags:\n" + + "\t- {}\n".format(flag0) + + "\t- {}\n".format(flag1) + + "in requests.get({}).text".format(total_url)) + raise Exception(msg) + out = out[ind0[0] + 1].split('') + nresults = len(out) - 1 + + # Get columns + col = [kk.replace('', '').replace('', '').strip() + for kk in out[0].split('')] + ncol = len(col) + colex = ['Wavelength', 'Ion', 'Data Type', 'Transition', 'File Details'] + if col != colex: + msg = ("Format of table columns in html seems to have changed!\n" + + "\t- expected: {}\n".format(colex) + + "\t- observed: {}".format(col)) + raise Exception(msg) + + lout = [] + for ii in range(0, nresults): + lstri = out[ii+1].split('') + assert len(lstri) == ncol + lamb = lstri[0].replace('Å', '') + elm, charg = lstri[1].replace('
', '').split('') + if charg == '+': + charg = '1+' + if element is not None and elm.lower() != element.lower(): + continue + if charge is not None and charg != charge: + continue + typ = lstri[2].replace('', '').split('>')[1] + trans = lstri[3].replace(' ', ' ') + trans = trans.replace('', '^{').replace('', '}') + trans = trans.replace('', '_{').replace('', '}') + trans = trans.replace('→', '->') + fil = lstri[4][lstri[4].index('detail')+len('detail'):] + fil = fil[:fil.index('.dat')+len('.dat')] + lout.append([lamb, elm, charg, typ, trans, fil]) + + # Format output + char = np.array(lout) + col = ['Wavelength', 'Element', 'Charge', 'Data Type', + 'Transition', 'Full file name'] + arr = _getcharray(char, col=col, + sep=' ', line='-', just='l', + returnas=returnas, verb=verb) + return arr + + +# ############################################################################# +# Download +# ############################################################################# + + +def _check_exists(filename, update=None): + + # In case a small modification becomes necessary later + target = filename + path_local = _get_PATH_LOCAL() + + # Check whether the local .tofu repo exists, if not recommend tofu-custom + if path_local is None: + path = os.path.join(os.path.expanduser('~'), '.tofu', 'openadas2tofu') + msg = ("You do not seem to have a local ./tofu repository\n" + + "tofu uses that local repository to store all user-specific " + + "data and downloads\n" + + "In particular, openadas files are downloaded and saved in:\n" + + "\t{}\n".format(path) + + " => to set-up your local .tofu repo, run in a terminal:\n" + + "\ttofu-custom") + raise Exception(msg) + + # Parse intermediate repos and create if necessary + lrep = target.split('/')[1:-1] + for ii in range(len(lrep)): + repo = os.path.join(path_local, *lrep[:ii+1]) + if not os.path.isdir(repo): + os.mkdir(repo) + + # Check if file already exists + path = os.path.join(path_local, *lrep) + pfe = os.path.join(path, target.split('/')[-1]) + if os.path.isfile(pfe): + if update is False: + msg = ("File already exists in your local repo:\n" + + "\t{}\n".format(pfe) + + " => if you want to force download, use update=True" + + " (local file will be overwritten)") + raise FileAlreayExistsException(msg) + else: + return True, pfe + else: + return False, pfe + + +def download(filename=None, + update=None, verb=None, returnas=None): + """ Download desired file from https://open.adas.ac.uk + + All downloaded files are stored in your local tofu directory (~/.tofu/) + + example + ------- + >>> import tofu as tf + >>> filename = '/adf15/pec40][ar/pec40][ar_ls][ar16.dat' + >>> tf.openadas2tofu.download(filename) + """ + + # --------------------------- + # Check + if verb is None: + verb = True + if update is None: + update = False + if returnas is None: + returnas = False + + c0 = (not isinstance(filename, str) + or filename[:4] != '/adf' or filename[-4:] != '.dat') + if c0: + msg = ("filename must be a str (full file name) of the form:\n" + + "\t/adf.../.../....dat") + raise Exception(msg) + url = _URL_DOWNLOAD + filename + + exists, pfe = _check_exists(filename, update=update) + if exists is True and verb is True: + msg = ("File already exists, will be downloaded and overwritten:\n" + + "\t{}".format(pfe)) + warnings.warn(msg) + + # --------------------------- + # Download + # Note the stream=True parameter below + with requests.get(url, stream=True) as rr: + rr.raise_for_status() + with open(pfe, 'wb') as ff: + for chunk in rr.iter_content(chunk_size=8192): + # filter-out keep-alive new chunks + if chunk: + ff.write(chunk) + # ff.flush() + + if verb is True: + msg = ("file {} was copied to:\n".format(filename) + + "\t{}".format(pfe)) + print(msg) + if returnas is str: + return pfe + + +def download_all(files=None, searchstr=None, + lambmin=None, lambmax=None, element=None, + include_partial=None, update=None, verb=None): + """ Download all desired files from https://open.adas.ac.uk + + The files to download can be provided either as: + - a list of full openadas file names (files) + - the result of an online freeform search + (searchstr fed to search_online) + - the result of an online search by wavelength + (lambmin, lambmax and element fed to search_online_by_wavelengthA) + + All downloaded files are stored in your local tofu directory (~/.tofu/) + + example + ------- + >>> import tofu as tf + >>> tf.openadas2tofu.download_all(lambmin=3., lambmax=4., element='ar') + """ + # Check + if include_partial is None: + include_partial = False + if update is None: + update = False + if verb is None: + verb = True + lc = [files is not None, + searchstr is not None, + any([ss is not None for ss in [lambmin, lambmax, element]])] + if np.sum(lc) != 1: + msg = ( + "Please either searchstr xor (lambmin, lambmax, element)\n" + + "\t- files: list of full file names to be downloaded\n" + + "\t- searchstr: uses search_online()\n" + + "\t- lambmin, lambmax, element: search_online_by_wavelengthA") + raise Exception(msg) + + # Get list of files + if lc[0]: + if isinstance(files, str): + files = [files] + if not (isinstance(files, list) + and all([isinstance(ss, str) for ss in files])): + msg = "files must be a list of full openadas file names!" + raise Exception(msg) + elif lc[1]: + arr = search_online(searchstr=searchstr, + include_partial=include_partial, + verb=False, returnas=np.ndarray) + files = arr[:, -1] + elif lc[2]: + arr = search_online_by_wavelengthA(lambmin=lambmin, + lambmax=lambmax, + element=element, verb=False, + returnas=np.ndarray) + files = np.unique(arr[:, -1]) + + # Download + if verb is True: + msg = "Downloading from {} into {}:".format(_URL, _get_PATH_LOCAL()) + print(msg) + for ii in range(len(files)): + try: + exists, pfe = _check_exists(files[ii], update=update) + pfe = download(filename=files[ii], update=update, + verb=False, returnas=str) + if exists is True: + msg = "\toverwritten: \t{}".format(files[ii]) + else: + msg = "\tdownloaded: \t{}".format(files[ii]) + except FileAlreayExistsException: + msg = "\talready exists: {}".format(files[ii]) + except Exception as err: + msg = (str(err) + + "\n\nCould not download file {}".format(files[ii])) + raise err + + if verb is True: + print(msg) + + +def clean_downloads(): + """ Delete all openadas files downloaded in your ~/.tofu/ directory """ + path_local = _get_PATH_LOCAL() + if path_local is None: + return + lf = [ff for ff in os.listdir(path_local) + if os.path.isfile(os.path.join(path_local, ff))] + ld = [ff for ff in os.listdir(path_local) + if os.path.isdir(os.path.join(path_local, ff))] + for ff in lf: + os.remove(ff) + for dd in ld: + shutil.rmtree(os.path.join(path_local, dd)) diff --git a/tofu/scripts/tofucustom.py b/tofu/scripts/tofucustom.py index 4df51a866..04201c98b 100755 --- a/tofu/scripts/tofucustom.py +++ b/tofu/scripts/tofucustom.py @@ -12,11 +12,13 @@ # default values ################################################### + _SOURCE = os.path.abspath(os.path.join(os.path.dirname(__file__), '..')) _USER = getpass.getuser() _USER_HOME = os.path.expanduser('~') _TARGET = os.path.join(_USER_HOME, '.tofu') _LF = ['_imas2tofu_def.py', '_scripts_def.py'] +_LD = ['openadas2tofu'] ################################################### @@ -24,7 +26,9 @@ # function ################################################### -def custom(target=_TARGET, source=_SOURCE, files=_LF): + +def custom(target=_TARGET, source=_SOURCE, + files=_LF, directories=_LD): # Caveat (up to now only relevant for _TARGET) if target != _TARGET: @@ -40,9 +44,16 @@ def custom(target=_TARGET, source=_SOURCE, files=_LF): # Try creating directory and copying modules try: + # Create .tofu/ if non-existent if not os.path.isdir(target): os.mkdir(target) + # Create directories + for dd in directories: + if not os.path.isdir(os.path.join(target, dd)): + os.mkdir(os.path.join(target, dd)) + + # Copy files for ff in files: mod, f0 = ff.split('_')[1:] copyfile(os.path.join(source, mod, '_'+f0), diff --git a/tofu/version.py b/tofu/version.py index 7ece8f61e..443c713ec 100644 --- a/tofu/version.py +++ b/tofu/version.py @@ -1,2 +1,2 @@ # Do not edit, pipeline versioning governed by git tags! -__version__ = '1.4.2b4-71-g44a9874f' +__version__ = '1.4.2b4-92-gf832ab7e'