# -*- coding: utf-8 -*-
"""
@author: jeremy leconte
Library of useful functions for users (only).
Functions here CANNOT be called into the library itself: importing this module in others would
lead to recursive import problems.
"""
import os
import pickle
import numpy as np
from exo_k.ktable import Ktable
from exo_k.ktable5d import Ktable5d
from exo_k.xtable import Xtable
from exo_k.hires_spectrum import Hires_spectrum
from exo_k.util.cst import PI
from exo_k.atm_evolution.atm_evol import Atm_evolution
from exo_k.aerosol import mmr_to_number_density # noqa
[docs]
def create_table(logpgrid=None, tgrid=None, xgrid=None, wngrid=None, wnedges=None,
ggrid=None, weights=None, p_unit='Pa', kdata_unit='m^2/molecule',
value=None, mol=None):
"""Create a new table with the required dimensions padded with zeros.
The class of the object created will depend on the dimensions required.
Parameters
----------
logpgrid: array, np.ndarray
Grid in log p
tgrid: array, np.ndarray
Grid in T
xgrid: array, np.ndarray
Grid in volume mixing ratio
wngrid: array, np.ndarray
Grid wavenumber
ggrid: array, np.ndarray
Grid g-space
weights: array, np.ndarray
Grid of quadrature weights
p_unit: str (optional)
Unit of pressure (default Pa)
kdata_unit: str (optional)
Unit of cross section (default is m^2/molecule)
Returns
-------
new_table: `Xtable`, `Ktable`, or `Ktable5d`
A new table
"""
if (ggrid is None) != (weights is None):
raise RuntimeError("Both ggrid and weights should be specified or equal to None.")
if (logpgrid is None) or (tgrid is None) or (wngrid is None):
raise RuntimeError("You should at least specify logpgrid, tgrid, and wngrid.")
if ggrid is not None:
if xgrid is not None:
new_table=Ktable5d()
new_table.xgrid=np.array(xgrid, dtype=float)
new_table.Nx=new_table.xgrid.size
else:
new_table=Ktable()
new_table.ggrid=np.array(ggrid, dtype=float)
new_table.weights=np.array(weights, dtype=float)
if new_table.ggrid.size != new_table.weights.size:
raise RuntimeError('ggrid and weights do not have the same length')
new_table.Ng=new_table.ggrid.size
else:
new_table=Xtable()
new_table.logpgrid=np.array(logpgrid, dtype=float)
new_table.pgrid=10**new_table.logpgrid
new_table.tgrid=np.array(tgrid, dtype=float)
new_table.wns=np.array(wngrid, dtype=float)
if wnedges is not None:
new_table.wnedges=np.array(wnedges, dtype=float)
else:
new_table.wnedges=np.concatenate( \
([new_table.wns[0]],(new_table.wns[:-1]+new_table.wns[1:])*0.5,[new_table.wns[-1]]))
new_table.Np=new_table.pgrid.size
new_table.Nt=new_table.tgrid.size
new_table.Nw=new_table.wns.size
sh=[new_table.Np,new_table.Nt,new_table.Nx,new_table.Nw,new_table.Ng]
sh=[x for x in sh if x is not None]
if value is None:
new_table.kdata=np.zeros(sh, dtype=float)
else:
new_table.kdata=np.full(sh, value, dtype=float)
new_table.p_unit=p_unit
new_table.kdata_unit=kdata_unit
new_table.mol=mol
return new_table
[docs]
def hires_to_ktable(filename_grid=None, xgrid=None, **kwargs):
"""Emulates :func:`exo_k.ktable.Ktable.hires_to_ktable`
and :func:`exo_k.ktable5d.Ktable5d.hires_to_ktable` as functions and not methods
so that the user can call them without creating a Ktable first.
See those methods for details on the available arguments and options.
"""
if filename_grid is not None:
if len(np.array(filename_grid).shape) == 3 and xgrid is None:
raise RuntimeError("""From the shape of filename_grid
it seems you want have an xgrid dimension, but xgrid is None!
""")
if xgrid is not None: # creates a Ktable5d
if filename_grid is not None:
if len(np.array(filename_grid).shape) != 3:
raise RuntimeError("""You provided a xgrid of vmrs,
but the shape of filename_grid is not
of the type (Np, Nt, Nx)!
""")
res=Ktable5d()
res.hires_to_ktable(filename_grid=filename_grid, xgrid=xgrid, **kwargs)
else:
res=Ktable()
res.hires_to_ktable(filename_grid=filename_grid, **kwargs)
return res
[docs]
def hires_to_xtable(**kwargs):
"""Emulates :func:`exo_k.xtable.Xtable.hires_to_xtable`
as a function and not a method
so that the user can call it without creating an Xtable first.
See :func:`~exo_k.xtable.Xtable.hires_to_xtable` for details
on the available arguments and options.
"""
res=Xtable()
res.hires_to_xtable(**kwargs)
return res
[docs]
def convert_kspectrum_to_hdf5(file_in, file_out=None, **kwargs):
"""Converts kspectrum like spectra to hdf5 format for speed and space.
Helper function. Real work done in :class:`~exo_k.util.hires_spectrum.Hires_spectrum`
__init__ funtion.
Go there to see all the available arguments and options.
Parameters
----------
file_in: str
Initial kspectrum filename.
file_out: str
Name of the final hdf5 file to be created. If not provided,
'file_in.h5' will be used.
"""
tmp=Hires_spectrum(file_in, **kwargs)
if file_out is None: file_out=file_in
tmp.write_hdf5(file_out)
[docs]
def create_fname_grid(base_string, logpgrid=None, tgrid=None, xgrid=None,
p_kw=None, t_kw=None, x_kw=None):
"""Creates a grid of filenames from an array of pressures, temperatures (
and vmr if there is a variable gas).
Parameters
----------
base_string: str
Generic name of the spectra files with specific keywords to be replaced
by the relevant numerical values
logpgrid: Array
Grid in log(pressure/Pa) of the input
tgrid: Array
Grid in temperature of the input
xgrid: Array
Input grid in vmr of the variable gas
.. warning::
The result of this function is much more predictable
if the values in the above arrays are given as integers or directly strings.
If you want to use floats anyway, good luck.
Parameters
----------
p_kw: str
t_kw: str
x_kw: str
The pattern string that will be recognized as keywords between
{} in base_string (See examples).
Examples
--------
>>> logpgrid=[1,2]
>>> tgrid=np.array([100.,150.,200.])
>>> file_grid=exo_k.create_fname_grid('spectrum_CO2_1e{logp}Pa_{t}K.hdf5',
logpgrid=logpgrid,tgrid=tgrid,p_kw='logp',t_kw='t')
array([['spectrum_CO2_1e1Pa_100K.hdf5', 'spectrum_CO2_1e1Pa_150K.hdf5',
'spectrum_CO2_1e1Pa_200K.hdf5'],
['spectrum_CO2_1e2Pa_100K.hdf5', 'spectrum_CO2_1e2Pa_150K.hdf5',
'spectrum_CO2_1e2Pa_200K.hdf5']], dtype='<U28')
"""
logpgrid=np.array(logpgrid)
tgrid=np.array(tgrid)
res=[]
if xgrid is None:
for iP in range(logpgrid.size):
for iT in range(tgrid.size):
dict_opt={p_kw:str(logpgrid[iP]),t_kw:str(tgrid[iT])}
fname=base_string.format(**dict_opt)
res.append(fname)
return np.array(res).reshape((logpgrid.size,tgrid.size))
else:
xgrid=np.array(xgrid)
for iP in range(logpgrid.size):
for iT in range(tgrid.size):
for iX in range(xgrid.size):
dict_opt={p_kw:str(logpgrid[iP]), \
t_kw:str(tgrid[iT]),x_kw:str(xgrid[iX])}
fname=base_string.format(**dict_opt)
res.append(fname)
return np.array(res).reshape((logpgrid.size,tgrid.size,xgrid.size))
[docs]
def finalize_LMDZ_dir(corrkname, IRsize, VIsize):
"""Creates the right links for a LMDZ type directory to be read by the LMDZ generic GCM.
You will need to create a proper Q.dat before using with the LMDZ GCM.
.. important::
You must have run :func:`exo_k.ktable.Ktable.write_LMDZ` or
:func:`exo_k.ktable5d.Ktable5d.write_LMDZ` for both of your IR and VI channels
beforehand.
Parameters
----------
corrkname: str
Path to the directory with the LMDZ ktable to finalize
IRsize: int
Number of IR spectral bins
VIsize: int
Number of VI spectral bins
"""
newdir=os.path.join(corrkname,str(IRsize)+'x'+str(VIsize))
try:
os.mkdir(newdir)
except FileExistsError:
os.system('rm -rf '+newdir)
os.mkdir(newdir)
os.symlink('../IR'+str(IRsize)+'/corrk_gcm_IR.dat',os.path.join(newdir,'corrk_gcm_IR.dat'))
os.symlink('../IR'+str(IRsize)+'/narrowbands_IR.in',os.path.join(newdir,'narrowbands_IR.in'))
os.symlink('../VI'+str(VIsize)+'/corrk_gcm_VI.dat',os.path.join(newdir,'corrk_gcm_VI.dat'))
os.symlink('../VI'+str(VIsize)+'/narrowbands_VI.in',os.path.join(newdir,'narrowbands_VI.in'))
print('Everything went ok. Your ktable is in:',newdir)
print("You'll probably need to add Q.dat before using it though!")
[docs]
def load_atm_evolution(pickle_filename):
"""Load an Atm_evolution instance from a pickle file
Parameters
----------
pickle_filename: str
name of pickle file
"""
with open(pickle_filename, 'rb') as filehandler:
atm_evol = pickle.load(filehandler)
return atm_evol