from importlib import util as imp_u, import_module as imp
from os import listdir
from mango.constants import c
[docs]def read(save_type):
"""
Import modules based on options for reading data files.
Parameters
----------
save_type: str
file type extension desired
Returns
-------
load: func
loading files function
save_type: str
save type based on available
"""
if save_type == "hdf5":
try:
from mango.hdf5io import load
new_save_type = 'hdf5'
except ImportError:
from pickle import load
c.Error("M Using pickle for data reading\n install pytables for hdf5 r/w")
new_save_type = "pkl"
elif save_type == "pkl":
from pickle import load
new_save_type = "pkl"
elif save_type == 'xyz':
load = None
new_save_type = 'xyz'
else:
new_save_type = "txt"
load = None
if save_type != new_save_type:
c.Error(f'W Unable to load {save_type} reader')
return load, save_type
[docs]def write(save_type):
"""
Import modules based on options for writing data files.
Parameters
----------
save_type: str
file type extension desired
Returns
-------
dump: func
save files function
save_type: str
save type based on available
"""
if save_type == "hdf5":
try:
from mango.hdf5io import save as dump
except ImportError:
from pickle import dump
c.Error("M Using pickle for data reading\n install pytables for hdf5 r/w")
save_type = "pkl"
elif save_type == "pkl":
from pickle import dump
else:
save_type = "txt"
dump = None
return dump, save_type
[docs]def inquirer():
"""
Import inquirer if required.
Returns
-------
bool:
Is inquirere available?
tuple:
inquirer functions
"""
if c.opsys != "nt" and c.tinfo['otty']:
try:
from mango.pp.util import captured_output
with captured_output() as (out, err):
from inquirer import Checkbox
from inquirer.render.console import ConsoleRender
return True, (Checkbox, ConsoleRender)
except ImportError:
return False, (None, None)
[docs]class matplotlib():
"""Importing matplotlib when necessary, allows for changing the backend before import."""
def __init__(self, backend=None):
from matplotlib import use, rcParams
if c.havedisplay is False:
use('Agg')
rcParams['agg.path.chunksize'] = 10000
elif backend is not None:
use(backend)
from matplotlib.pyplot import figure, savefig
self.figure = figure
self.savefig = savefig
self.rcParams = rcParams
[docs] def suscep(self):
from matplotlib.lines import Line2D
from matplotlib.pyplot import legend, yscale
return (self.figure, self.savefig, self.errorbarpicker(),
Line2D, legend, yscale)
[docs] def col(self):
return self.figure, self.savefig, self.htuple()
[docs] @staticmethod
def showg():
from matplotlib.pyplot import show, get_fignums
return show, get_fignums
[docs] @staticmethod
def close():
from matplotlib.pyplot import close
return close
[docs] @staticmethod
def errorbarpicker():
import matplotlib.legend_handler as mlh
from matplotlib.container import ErrorbarContainer
class newHandler(mlh.HandlerErrorbar):
def create_artists(self, *args, **kwargs):
a_list = mlh.HandlerErrorbar.create_artists(self, *args, **kwargs)
a_list = a_list[-2:] + a_list[:-2]
return a_list
return {ErrorbarContainer: newHandler()}
[docs] def prettyplot(self):
from matplotlib import ticker, __version__ as mplversion
from matplotlib.pyplot import setp, gcf
from matplotlib.backends.backend_pdf import PdfPages
from subprocess import getstatusoutput
if getstatusoutput('which latex')[0] == 0:
self.rcParams['text.usetex'] = True
self.rcParams['font.family'] = 'MathJax_Main'
self.rcParams['font.size'] = 14
if int(mplversion[0]) < 3:
self.rcParams['text.latex.unicode'] = True
return self.figure, PdfPages, setp, gcf, ticker, self.htuple()
[docs] def htuple(self):
"""single legend entry for multiple connected graphs with one legend. see mango.tools.plotting"""
import matplotlib.legend_handler as lh
import numpy as np
from matplotlib.collections import PolyCollection
class mytuple(lh.HandlerTuple):
def create_artists(self, legend, orig_handle,
xdescent, ydescent, width, height, fontsize,
trans):
handler_map = legend.get_legend_handler_map()
if self._ndivide is None:
ndivide = len(orig_handle)
else:
ndivide = self._ndivide
if self._pad is None:
pad = legend.borderpad * fontsize
else:
pad = self._pad * fontsize
if ndivide > 1:
width = (width - pad * (ndivide - 1)) / ndivide
xds_cycle = lh.cycle(xdescent - (width + pad) * np.arange(ndivide))
a_list = []
c_list = [0, 1] if isinstance(orig_handle[0], PolyCollection) else [0]
for no, handle1 in enumerate(orig_handle):
handler = legend.get_legend_handler(handler_map, handle1)
if no in c_list:
_a_list = handler.create_artists(
legend, handle1,
next(xds_cycle), ydescent, width, height, fontsize, trans)
a_list.extend(_a_list)
return a_list
return {tuple: mytuple()}
[docs]def fftw():
"""
Import the fastest version of fftw installed.
Returns
-------
output: fft, ffts, ifft
Fastest fft implementations available
"""
from scipy import fftpack as ffts
try:
from pyfftw.interfaces import scipy_fftpack
from pyfftw import byte_align
fft = scipy_fftpack.fft
ifft = scipy_fftpack.ifft
except ImportError:
def fft(arr, n=None, axis=-1, threads=None):
return ffts.fft(arr, n=n, axis=axis)
def ifft(arr, n=None, axis=-1, threads=None):
return ffts.ifft(arr, n=n, axis=axis)
def byte_align(array, n=None, dtype=None):
return array
c.Error(">M Using scipy fft for Susceptibility\n install fftw and pyfftw for better performance")
return fft, ffts, ifft, byte_align
[docs]def getpotentials(name):
"""
Load pluggable potentials.
Parameters
----------
name: string
Filename of potentials module to import
Returns
-------
output: module
Loaded module
"""
try:
module = imp("mango.potentials.{}".format(name))
except ImportError:
module = _loadfromspec(name)
c.Error(f"M Loaded {module.__name__}")
try:
for i in c.posfuncs:
if i not in module.__all__:
c.Error("W function {} not replaced.".format(i))
except AttributeError:
c.Error("F module has no attribute __all__")
return module
def _loadfromspec(name):
full_dirname = name.rsplit("/", 1)
dirname = '.' if len(full_dirname[0]) == 0 else full_dirname[0]
try:
mod_files = [file for file in listdir(dirname) if file.startswith(full_dirname[1])]
except FileNotFoundError:
c.Error("F Can't find module file '{}'".format(name))
requested = full_dirname[1] if full_dirname[1] in mod_files else mod_files[0]
if len(mod_files) > 1:
c.Error("{}{}".format("W Multiple files start with '{}'\n".format(full_dirname[1]),
"Assuming module is '{}'".format(requested)))
mod_file = f"{dirname}/{requested}"
try:
spec = imp_u.spec_from_file_location(mod_file.rsplit("/")[-1].split(".")[0], mod_file)
module = imp_u.module_from_spec(spec)
spec.loader.exec_module(module)
except (AttributeError, ImportError):
c.Error("F File '{}' is not a module".format(mod_files[0]))
return module