def choose_pool(mpi=False, processes=1, **kwargs):
"""
Extends the capabilities of the schwimmbad.choose_pool method.
It handles the `use_dill` parameters in kwargs, that would otherwise raise an error when processes > 1.
Any thread in the returned multiprocessing pool (e.g. processes > 1) also default
Docstring from schwimmbad:
mpi : bool, optional
Use the MPI processing pool, :class:`~schwimmbad.mpi.MPIPool`. By
default, ``False``, will use the :class:`~schwimmbad.serial.SerialPool`.
processes : int, optional
Use the multiprocessing pool,
:class:`~schwimmbad.multiprocessing.MultiPool`, with this number of
processes. By default, ``processes=1``, will use the
:class:`~schwimmbad.serial.SerialPool`.
**kwargs
Any additional kwargs are passed in to the pool class initializer selected by the arguments.
"""
if processes == 1 or mpi:
pool = schwimmbad.choose_pool(mpi=mpi, processes=1, **kwargs)
is_master = pool.is_master()
else:
if 'use_dill' in kwargs:
# schwimmbad MultiPool does not support dill so we remove this option from the kwargs
_ = kwargs.pop('use_dill')
pool = schwimmbad.choose_pool(mpi=False, processes=processes, **kwargs)
# this MultiPool has no is_master() attribute like the SerialPool and MpiPool
# all threads will then be 'master'.
is_master = True
return pool, is_master
def choose_pool(processes, mpi=False):
""" Get processing pool
"""
do_mpi, size, rank = use_mpi(require_mpi=mpi)
if do_mpi:
try:
import schwimmbad
pool = schwimmbad.choose_pool(mpi=do_mpi, processes=(size - 1))
pool.broadcast = types.MethodType(_dummy_broadcast, pool)
atexit.register(pool.close)
if processes:
logging.info('NOTE: that for MPI process size determined by '
'MPI launch size, not the processes argument')
if do_mpi and not mpi:
logging.info('NOTE: using MPI as this process was launched'
'under MPI')
except ImportError:
raise ValueError("Failed to start up an MPI pool, "
"install mpi4py / schwimmbad")
elif processes == 1:
pool = SinglePool()
else:
pool = BroadcastPool(processes)
pool.size = processes
if size:
pool.size = size
return pool
def run_mcmc(self,
params,
pids,
nwalkers=100,
nsamples=100,
mpi=False,
threads=1):
"""
Runs an MCMC chain
"""
# Setup the pool, to map lnprob
import schwimmbad
import emcee
pool = schwimmbad.choose_pool(mpi=mpi, processes=threads)
ndim = len(params)
pos0 = [
params + params * np.random.randn(ndim) * 0.01
for i in range(nwalkers)
]
sampler = emcee.EnsembleSampler(nwalkers,
ndim,
self,
args=(pids, ),
pool=pool)
sampler.run_mcmc(pos0, nsamples)
pool.close()
return sampler.chain
def compute_coeff_parallel(pos, nmax, lmax, r_s, ncores):
"""
"""
#pool = schwimmbad.choose_pool(mpi=False, processes=ncores)
pool = schwimmbad.choose_pool(mpi=True, processes=2)
results = main(pool, pos, mass, nmax, lmax, r_s)
Snlm, Tnlm = coeff_matrix(results)
return Snlm, Tnlm
def main(n_tasks, log_pathname):
"""
Simulate the execution of a bunch of tasks using MPI,
as well as setup MPI logging.
We should see all processor ranks logging similar events right up
till `pool = schwimmbad.choose_pool(mpi=True)` is called.
After which only rank 0 will be logging events, and executing lines
of code.
All tasks should contain the same value for rank as an argument,
and the value will be zero.
The tasks themeselves will be carried out by the workers,
ranks 1 and higher, and rank 0 will act as a scheduler.
"""
# comm and processor info
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
stream = MPIStreamIO(log_pathname)
logger_factory = MPILoggerFactory(stream)
structlog.configure(
processors=DEFAULT_PROCESSORS,
logger_factory=logger_factory,
)
# bind the logs to include the rank and size
log = structlog.get_logger('status').bind(size=size)
log.info('checking rank and size')
# define a pool which tells workers to wait for the scheduler's instructions
log.info('about to initialise an MPI pool')
pool = schwimmbad.choose_pool(mpi=True)
# after pool is created, only rank 0 will be executing the following lines
log.info('MPI pool chosen')
task_ids = range(1, n_tasks + 1)
# the task class that calls some function
task_calculator = SomeTask()
log.info('task chosen')
# define the list of tasks and the input parameters
args = [(task_id, rank) for task_id in task_ids]
# the rank argument should always be zero as we defined an MPI pool above
log.info(args=args)
# map tasks across the pool of workers
# rank 0 acts only as a scheduler, ranks 1 and higher are workers
pool.map(task_calculator, args)
pool.close()
log.info('finished processing')
def choose_pool(processes, mpi=False):
if mpi:
try:
import schwimmbad
pool = schwimmbad.choose_pool(mpi=mpi, processes=processes)
pool.broadcast = types.MethodType(_dummy_broadcast, pool)
except ImportError:
raise ValueError("Failed to start up an MPI pool, "
"install mpi4py / schwimmbadd")
else:
pool = BroadcastPool(processes)
return pool
def choose_pool(processes, mpi=False):
if mpi:
try:
import schwimmbad
pool = schwimmbad.choose_pool(mpi=mpi,
processes=processes)
pool.broadcast = types.MethodType(_dummy_broadcast, pool)
except ImportError:
raise ValueError("Failed to start up an MPI pool, "
"install mpi4py / schwimmbadd")
else:
pool = BroadcastPool(processes)
return pool
def test_ee_g_pool(self):
from functools import partial
import numpy as np
import schwimmbad
from pyeee import ee
from pyeee.functions import G
from pyeee.utils import func_wrapper
# Function and parameters
func = G
npars = 6
params = [78., 12., 0.5, 2., 97., 33.] # G
# Partialise function with fixed parameters
arg = [params]
kwarg = {}
obj = partial(func_wrapper, func, arg, kwarg)
# Screening
lb = np.zeros(npars)
ub = np.ones(npars)
nprocs = 4
ipool = schwimmbad.choose_pool(mpi=False, processes=nprocs)
out = ee(obj,
lb,
ub,
self.nt,
x0=None,
mask=None,
ntotal=self.ntotal,
nsteps=self.nsteps,
processes=nprocs,
pool=ipool)
ipool.close()
# Check
self.assertEqual(list(np.around(out[:, 0], 3)),
[0.045, 0.24, 1.624, 0.853, 0.031, 0.084])
def test_eee_k(self):
from functools import partial
import os
import numpy as np
import schwimmbad
from pyeee import eee
from pyeee.utils import func_wrapper
from pyeee.functions import bratley
# Function and parameters
func = bratley
npars = 10
params = [] # k
# Screening
lb = np.zeros(npars)
ub = np.ones(npars)
nprocs = 4
ipool = schwimmbad.choose_pool(mpi=False, processes=nprocs)
out = eee(func,
lb,
ub,
mask=None,
ntfirst=self.ntfirst,
ntlast=self.ntlast,
nsteps=self.nsteps,
processes=nprocs,
pool=ipool,
logfile='tlog.txt')
ipool.close()
# Check
self.assertEqual(list(np.where(out)[0] + 1),
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
self.assertTrue(os.path.exists('tlog.txt'))
# Clean
if os.path.exists('tlog.txt'): os.remove('tlog.txt')
def screening(func,
x0,
lb,
ub,
mask=None,
arg=(),
kwarg={},
nt=-1,
nsteps=6,
ntotal=-1,
processes=1,
pool=None,
verbose=0,
parameterfile=None,
parameterwriter=None,
objectivefile=None,
objectivereader=None,
shell=False,
debug=False):
"""
Morris Screening
Definition
----------
def screening(func, x0, lb, ub, mask=None,
arg=(), kwarg={},
nt=-1, nsteps=6, ntotal=-1,
processes=1,
verbose=0,
parameterfile=None, parameterwriter=None,
objectivefile=None, objectivereader=None,
shell=False, debug=False):
Input
-----
func python function or string for external executable
The objective function to screen.
x0 1D-array
Will be taken at dimensions with mask==False.
lb 1D-array
Lower bounds of parameters.
ub 1D-array
Upper bounds of parameters.
Optional Input
--------------
mask 1D-array (Default: include all parameters)
Include (1,True) or exclude (0,False) parameters in screening.
arg tuple (Default: empty tuple)
Additional arguments passed to objective and constraint functions (only for Python functions).
kwarg dict (Default: empty dict)
Additional keyword arguments passed to objective and constraint functions (only for Python functions).
nt int (Default: len(x0))
The number of trajectories to return.
ntotal int (Default: max(nt*10,nt**2))
Total number of trajectories for selecting optimal nt trajectories.
nsteps int (Default: 6)
Number of steps between lower and upper bounds.
processes int (Default: 1)
The number of processes to use to evaluate objective function and constraints.
pool schwimmbad pool object (Default: None)
Generic map function used from schwimmbad, which provides, serial, multiprocessor,
and MPI mapping functions. The pool is chosen with
schwimmbad.choose_pool(mpi=True/False, processes=nummultiprocessors).
The pool will be chosen automatically if pool is None.
MPI pools can only be opened and closed once. So if you want to use screening several
times in one program, then you have to choose the pool and later close the pool in the
wrapping progam and pass it to screening.
parameterfile string (Default: None)
Parameter file for executable; must be given if func is name of executable
parameterwriter function (Default: None)
Python function for writing parameter file if func is name of executable
objectivefile string (Default: None)
File with objective value from executable; must be given if func is name of executable
objectivereader function (Default: None)
Python function for reading objective value if func is name of executable
shell boolean (Default: False)
If True, the specified executable will be executed through the shell.
debug boolean (Default: False)
If True, model output is displayed for executable.
Output
------
2D-array - (nparameter,3)
if nt>1:
2D-array - (nparameter,3) with per parameter
1. mean of absolute elementary effects over all nt trajectories (mu*)
2. mean of elementary effects over all nt trajectories (mu)
3. standard deviation of elementary effects over all nt trajectories (sigma)
else:
2D-array - (nparameter,3) with per parameter
1. absolute elementary effect of each parameter
2. elementary effect of each parameter
2. zeros
License
-------
This file is part of the JAMS Python package, distributed under the MIT License.
Copyright (c) 2017 Matthias Cuntz - mc (at) macu (dot) de
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
History
-------
Written, Matthias Cuntz, Dec 2017
Modified, Matthias Cuntz, Dec 2017 - output for nt=1 also (npara,3)
Matthias Cuntz, Dec 2019 - bug: default ntotal was not set if ntotal<0 (but nt instead)
"""
warn('The function screening is deprecated from JAMS. Use module pyjams.',
category=DeprecationWarning)
# Get MPI communicator
try:
from mpi4py import MPI
comm = MPI.COMM_WORLD
csize = comm.Get_size()
crank = comm.Get_rank()
except ImportError:
comm = None
csize = 1
crank = 0
if csize > 1:
passrank = crank
else:
passrank = None
# Checks
# Function
assert hasattr(func, '__call__') or isinstance(
func, (str, list)), 'Invalid function handle or external call.'
# Bounds
assert len(x0) == len(
lb), 'Initial values and lower bounds must have the same lengths.'
assert len(lb) == len(
ub), 'Lower- and upper-bounds must have the same lengths.'
x0 = np.array(x0)
lb = np.array(lb)
ub = np.array(ub)
# Mask
if mask is None:
assert np.all(
ub >= lb
), 'All upper-bounds must be greater or equal than lower-bounds.'
else:
assert len(mask) == len(
ub), 'Mask and bounds must have the same lengths.'
if not np.all(mask):
assert len(mask) == len(
x0), 'Mask and x0 must have the same lengths.'
assert np.all(
ub[mask] >= lb[mask]
), 'All unmasked upper-bounds must be greater or equal than lower-bounds.'
# Parameterfile etc. keywords if func is name of executable
if isinstance(func, (str, list)):
if parameterfile is None:
raise IOError(
'parameterfile must be given if func is name of executable.')
if parameterwriter is None:
raise IOError(
'parameterwrite must be given if func is name of executable.')
if objectivefile is None:
raise IOError(
'objectivefile must be given if func is name of executable.')
if objectivereader is None:
raise IOError(
'objectivereader must be given if func is name of executable.')
# Set defaults
npara = len(lb)
if mask is None:
imask = np.ones(npara, dtype=np.bool)
else:
imask = mask
nmask = np.sum(imask)
if nt <= 0:
nt = npara
if ntotal <= 0:
ntotal = max(nt**2, 10 * nt)
# Partialise objective function
if isinstance(func, (str, list)):
obj = partial(_ext_obj_wrapper, func, lb, ub, imask, parameterfile,
parameterwriter, objectivefile, objectivereader, shell,
debug, passrank)
else:
obj = partial(_obj_wrapper, func, arg, kwarg)
# Sample trajectories
if (crank == 0) and (verbose > 0): print('Sample trajectories')
tmatrix, tvec = morris_sampling(nmask,
lb[imask],
ub[imask],
N=ntotal,
p=nsteps,
r=nt,
Diagnostic=False)
# Set input vector to trajectories and masked elements = x0
x = np.tile(x0, tvec.size).reshape(tvec.size, npara) # default to x0
x[:, imask] = tmatrix # replaced unmasked with trajectorie values
# Choose the right mapping function: single, multiprocessor or mpi
if pool is None:
import schwimmbad
ipool = schwimmbad.choose_pool(mpi=False if csize == 1 else True,
processes=processes)
else:
ipool = pool
# Calculate all model runs
if (crank == 0) and (verbose > 0): print('Calculate objective functions')
fx = np.array(ipool.map(obj, x))
if pool is None: ipool.close()
# Calc elementary effects
if (crank == 0) and (verbose > 0): print('Calculate elementary effects')
sa, res = elementary_effects(nmask,
tmatrix,
tvec,
fx,
p=nsteps,
Diagnostic=False)
# Output with zero for all masked parameters
if nt == 1:
out = np.zeros((npara, 3))
out[imask, 0] = np.abs(sa[:, 0])
out[imask, 1] = sa[:, 0]
else:
out = np.zeros((npara, 3))
out[imask, :] = res
return out
args = parser.parse_args()
# Set logger level based on verbose flags
if args.verbosity != 0:
if args.verbosity == 1:
logger.setLevel(logging.DEBUG)
else: # anything >= 2
logger.setLevel(1)
elif args.quietness != 0:
if args.quietness == 1:
logger.setLevel(logging.WARNING)
else: # anything >= 2
logger.setLevel(logging.ERROR)
else: # default
logger.setLevel(logging.INFO)
if args.seed is not None:
np.random.seed(args.seed)
pool = choose_pool(mpi=args.mpi, processes=args.n_cores)
logger.info("Using pool: {}".format(pool.__class__))
main(db_path=args.db_path,
data_root_path=args.data_root_path,
run_name=args.run_name,
filename=args.filename,
overwrite=args.overwrite,
pool=pool)
# -*- coding: utf-8 -*-
import os
import sys
from dolphin.processor import Processor
import time
import schwimmbad
pool = schwimmbad.choose_pool(mpi=True)
start_time = time.perf_counter()
cwd = os.getcwd()
base_path, _ = os.path.split(cwd)
processor = Processor(base_path)
lens_name = str(sys.argv[1])
model_id = str(sys.argv[2])
if pool.is_master():
print("Run [{}] for {} loaded.".format(model_id, lens_name))
processor.swim(lens_name, model_id=model_id)
if pool.is_master():
end_time = time.perf_counter()
print('Total time needed for computation: {:.2f} s'.format(end_time -
start_time))
def eee(func, *args, **kwargs):
"""
Parameter screening using Efficient/Sequential Elementary Effects of
Cuntz, Mai et al. (Water Res Research, 2015).
Note, the input function must be callable as `func(x)`.
Parameters
----------
func : callable
Python function callable as `func(x)` with `x` the function parameters.
lb : array_like
Lower bounds of parameters.
ub : array_like
Upper bounds of parameters.
x0 : array_like, optional
Parameter values used with `mask==0`.
mask : array_like, optional
Include (1,True) or exclude (0,False) parameters in screening (default: include all parameters).
ntfirst : int, optional
Number of trajectories in first step of sequential elementary effects (default: 5).
ntlast : int, optional
Number of trajectories in last step of sequential elementary effects (default: 5).
nsteps : int, optional
Number of intervals for each trajectory (default: 6)
weight : boolean, optional
If False, use the arithmetic mean mu* for each parameter if function has multiple outputs,
such as the mean mu* of each time step of a time series (default).
If True, return weighted mean mu*, weighted by sd.
seed : int or array_like
Seed for numpy``s random number generator (default: None).
processes : int, optinal
The number of processes to use to evaluate objective function and constraints (default: 1).
pool : `schwimmbad` pool object, optinal
Generic map function used from module `schwimmbad <https://schwimmbad.readthedocs.io/en/latest/>`_,
which provides, serial, multiprocessor, and MPI mapping functions (default: None).
The pool is chosen with:
schwimmbad.choose_pool(mpi=True/False, processes=processes).
The pool will be chosen automatically if pool is None.
MPI pools can only be opened and closed once. If you want to use screening several
times in one program, then you have to choose the pool, pass it to eee,
and later close the pool in the calling progam.
verbose : int, optional
Print progress report during execution if verbose>0 (default: 0).
logfile : File handle or logfile name
File name of possible log file (default: None = no logfile will be written).
plotfile : Plot file name
File name of possible plot file with fit of logistic function to mu* of first trajectories
(default: None = no plot produced).
Returns
-------
mask : ndarray
(len(lb),) Mask with 1=informative and 0=uninformative model parameters, to be used with '&' on input mask.
See Also
--------
:func:`~pyeee.screening.screening` : Elementary Effects, same as
:func:`~pyeee.screening.ee` : Elementary Effects
Examples
--------
>>> import numpy as np
>>> import pyeee
>>> seed = 1023
>>> np.random.seed(seed=seed)
>>> npars = 10
>>> lb = np.zeros(npars)
>>> ub = np.ones(npars)
>>> ntfirst = 10
>>> ntlast = 5
>>> nsteps = 6
>>> out = pyeee.eee(pyeee.K, lb, ub, x0=None, mask=None, ntfirst=ntfirst, ntlast=ntlast, nsteps=nsteps)
>>> print(np.where(out)[0] + 1)
[1 2 3 4 6]
History
-------
Written, Matthias Cuntz, Nov 2017
Modified, Matthias Cuntz, Jan 2018 - weight
Nov 2019 - plotfile, numpy docstring format
Matthias Cuntz, Jan 2020 - x0 optional
- verbose keyword
- distinguish iterable and array_like parameter types
Matthias Cuntz, Feb 2020 - ntsteps -> nsteps
- check if logfile is string instead of checking for file handle
"""
# Get keyword arguments
# This allows mixing keyword arguments of eee and keyword arguments to be passed to optimiser.
# The mixed syntax eee(func, *args, logfile=None, **kwargs) is only working in Python 3
# so need a workaround in Python 2, i.e. read all as keyword args and take out the keywords for eee.
x0 = kwargs.pop('x0', None)
mask = kwargs.pop('mask', None)
ntfirst = kwargs.pop('ntfirst', 5)
ntlast = kwargs.pop('ntlast', 5)
nsteps = kwargs.pop('nsteps', 6)
weight = kwargs.pop('weight', False)
seed = kwargs.pop('seed', None)
processes = kwargs.pop('processes', 1)
pool = kwargs.pop('pool', None)
verbose = kwargs.pop('verbose', 0)
logfile = kwargs.pop('logfile', None)
plotfile = kwargs.pop('plotfile', None)
# Set up MPI if available
try:
from mpi4py import MPI
comm = MPI.COMM_WORLD
csize = comm.Get_size()
crank = comm.Get_rank()
except ImportError:
comm = None
csize = 1
crank = 0
# Logfile
if crank == 0:
if logfile is None:
lfile = None
else:
# haswrite = getattr(logfile, "write", None)
# if haswrite is None:
# lfile = open(logfile, "w")
# else:
# if not callable(haswrite):
# lfile = logfile
# else:
# raise InputError('x0 must be given if mask is set')
if isinstance(logfile, str):
lfile = open(logfile, "w")
else:
lfile = logfile
else:
lfile = None
# Start
if crank == 0:
if (verbose > 0):
tee('Start screening in eee.', file=lfile)
else:
if lfile is not None:
print('Start screening in eee.', file=lfile)
# Check
assert len(args) == 2, 'lb and ub must be given as arguments.'
lb, ub = args[:2]
npara = len(lb)
if crank == 0:
assert len(lb) == len(ub), 'Lower and upper bounds have not the same size.'
lb = np.array(lb)
ub = np.array(ub)
# mask
if mask is None:
ix0 = np.ones(npara)
imask = np.ones(npara, dtype=np.bool)
iimask = np.arange(npara, dtype=np.int)
nmask = npara
else:
if x0 is None:
raise InputError('x0 must be given if mask is set')
ix0 = np.copy(x0)
imask = np.copy(mask)
iimask = np.where(imask)[0]
nmask = iimask.size
if nmask == 0:
if crank == 0:
if (verbose > 0):
tee('\nAll parameters masked, nothing to do.', file=lfile)
tee('Finished screening in eee.', file=lfile)
else:
if lfile is not None:
print('\nAll parameters masked, nothing to do.', file=lfile)
print('Finished screening in eee.', file=lfile)
if logfile is not None: lfile.close()
# Return all true
if mask is None:
return np.ones(len(lb), dtype=np.bool)
else:
return mask
if crank == 0:
if (verbose > 0):
tee('\nScreen unmasked parameters: ', nmask, iimask+1, file=lfile)
else:
if lfile is not None:
print('\nScreen unmasked parameters: ', nmask, iimask+1, file=lfile)
# Seed random number generator
if seed is not None: np.random.seed(seed=seed) # same on all ranks because trajectories are sampled on all ranks
# Choose the right mapping function: single, multiprocessor or mpi
if pool is None:
import schwimmbad
ipool = schwimmbad.choose_pool(mpi=False if csize == 1 else True, processes=processes)
else:
ipool = pool
# Step 1 of Cuntz et al. (2015) - first screening with ntfirst trajectories, calc mu*
res = screening( # returns (npara,3) with mu*, mu, std if nt>1
func, lb, ub, ntfirst,
x0=ix0, mask=imask,
nsteps=nsteps, ntotal=10*ntfirst,
processes=processes, pool=ipool,
verbose=0)
if res.ndim > 2:
if weight:
mustar = np.sum(res[:, iimask, 2] * res[:, iimask, 0],axis=0) / np.sum(res[:, iimask, 2], axis=0)
else:
mustar = np.mean(res[:, iimask, 0], axis=0)
else:
mustar = res[iimask, 0]
# Step 2 of Cuntz et al. (2015) - calc eta*
mumax = np.amax(mustar)
xx = np.arange(nmask) / np.float(nmask-1)
iisort = np.argsort(mustar)
yy = mustar[iisort] / mumax
if crank == 0:
if (verbose > 0):
tee('\nSorted means of absolute elementary effects (mu*): ', mustar[iisort], file=lfile)
tee('Normalised mu* = eta*: ', yy, file=lfile)
tee('Corresponding to parameters: ', iimask[iisort] + 1, file=lfile)
else:
if lfile is not None:
print('\nSorted means of absolute elementary effects (mu*): ', mustar[iisort], file=lfile)
print('Normalised mu* = eta*: ', yy, file=lfile)
print('Corresponding to parameters: ', iimask[iisort] + 1, file=lfile)
# Step 3.1 of Cuntz et al. (2015) - fit logistic function
# [y-max, steepness, inflection point, offset]
pini = np.array([yy.max(), (yy.max() - yy.max()) / xx.max(), 0.5 * xx.max(), yy.min()])
plogistic, f, d = opt.fmin_l_bfgs_b(cost_square,
pini,
args=(logistic_offset_p, xx, yy),
approx_grad=1,
bounds=[(None, None), (None, None), (None, None), (None, None)],
iprint=0,
disp=0)
# Step 3.2 of Cuntz et al. (2015) - determine point of steepest curvature -> eta*_thresh
def mcurvature(*args, **kwargs):
return -curvature(*args, **kwargs)
x_K = opt.brent(mcurvature, # x_K
args=(dlogistic, d2logistic, plogistic[0], plogistic[1], plogistic[2]),
brack=(xx[0], xx[-1]))
curvmax = logistic_offset_p(x_K, plogistic) # L(x_K)
eta_thresh = curvmax # eta*_thresh = L(x_K) # in range 0-1
if (curvmax > 0.2) or (x_K < xx[0]):
x_scaled = xx[0] # x_K = min(x)
eta_thresh = np.min(mustar) / mumax # eta*_thresh = min(mu*)/max(mu*)
mu_thresh = eta_thresh * mumax # mu*_thresh = eta*_thresh*max(mu*)
if crank == 0:
if (verbose > 0):
tee('\nThreshold eta*_thresh, mu*_tresh: ', eta_thresh, mu_thresh, file=lfile)
tee('L(x_K): ', logistic_offset_p(x_K, plogistic), file=lfile)
tee('p_opt of L: ', plogistic, file=lfile)
else:
if lfile is not None:
print('\nThreshold eta*_thresh, mu*_tresh: ', eta_thresh, mu_thresh, file=lfile)
print('L(x_K): ', logistic_offset_p(x_K, plogistic), file=lfile)
print('p_opt of L: ', plogistic, file=lfile)
# Plot first mu* of elementary effects with logistic function and threshold
if crank == 0:
if plotfile is not None:
try:
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
mpl.rcParams['font.family'] = 'sans-serif'
mpl.rcParams['font.sans-serif'] = 'Arial' # Arial, Verdana
mpl.rc('savefig', dpi=300, format='png')
if npara > 99:
mpl.rc('font', size=8)
else:
mpl.rc('font', size=11)
fig = plt.figure()
sub = plt.subplot(111)
xx = xx
yy = mustar[iisort]
line1 = sub.plot(xx, yy, 'ro')
nn = 1000
xx2 = xx.min() + np.arange(nn) / float(nn - 1) * (xx.max() - xx.min())
yy2 = logistic_offset_p(xx2, plogistic) * mumax
line2 = sub.plot(xx2, yy2, 'b-')
xmin, xmax = sub.get_xlim()
line3 = sub.plot([xmin, xmax], [mu_thresh, mu_thresh], 'k-')
if npara > 99:
xnames = ['{:03d}'.format(i) for i in iimask[iisort] + 1]
else:
xnames = ['{:02d}'.format(i) for i in iimask[iisort] + 1]
plt.setp(sub, xticks=xx, xticklabels=xnames)
plt.setp(sub, xlabel='Parameter')
plt.setp(sub, ylabel=r'$\mu*$')
fig.savefig(plotfile, transparent=False, bbox_inches='tight', pad_inches=0.035)
plt.close(fig)
except ImportError:
pass
# Step 4 of Cuntz et al. (2015) - Discard from next steps all parameters with
# eta* >= eta*_tresh, i.e. mu* >= mu*_tresh
imask[iimask] = imask[iimask] & (mustar < mu_thresh)
if np.all(~imask):
if crank == 0:
if (verbose > 0):
tee('\nNo more parameters to screen, i.e. all (unmasked) parameters are informative.', file=lfile)
tee('Finished screening in eee.', file=lfile)
else:
if lfile is not None:
print('\nNo more parameters to screen, i.e. all (unmasked) parameters are informative.', file=lfile)
print('Finished screening in eee.', file=lfile)
_cleanup(lfile, pool, ipool)
# Return all true
if mask is None:
return np.ones(len(lb), dtype=np.bool)
else:
return mask
# Step 5 and 6 of Cuntz et al. (2015) - Next trajectory with remaining parameters.
# Discard all parameters with |EE| >= mu*_tresh
# - Repeat until no |EE| >= mu*_tresh
niter = 1
donext = True
while donext:
if crank == 0:
if (verbose > 0):
tee('\nParameters remaining for iteration ', niter, ':', np.where(imask)[0] + 1, file=lfile)
else:
if lfile is not None:
print('\nParameters remaining for iteration ', niter, ':', np.where(imask)[0] + 1, file=lfile)
iimask = np.where(imask)[0]
res = screening( # returns EE(parameters) if nt=1
func, lb, ub, 1,
x0=ix0, mask=imask,
nsteps=nsteps, ntotal=10,
processes=processes, pool=ipool,
verbose=0)
# absolute EE
if res.ndim > 2:
if weight:
mustar = np.sum(res[:, iimask, 2] * res[:, iimask, 0], axis=0) / np.sum(res[:, iimask, 2], axis=0)
else:
mustar = np.mean(res[:, iimask, 0], axis=0)
else:
mustar = res[iimask, 0]
if crank == 0:
if (verbose > 0):
tee('Absolute elementary effects |EE|: ', mustar, file=lfile)
else:
if lfile is not None:
print('Absolute elementary effects |EE|: ', mustar, file=lfile)
imask[iimask] = imask[iimask] & (mustar < mu_thresh)
if np.all(~imask):
if crank == 0:
if (verbose > 0):
tee('\nNo more parameters to screen, i.e. all (unmasked) parameters are informative.', file=lfile)
tee('Finished screening in eee.', file=lfile)
else:
if lfile is not None:
print('\nNo more parameters to screen, i.e. all (unmasked) parameters are informative.', file=lfile)
print('Finished screening in eee.', file=lfile)
_cleanup(lfile, pool, ipool)
# Return all true
if mask is None:
return np.ones(len(lb), dtype=np.bool)
else:
return mask
# Step 6 of Cuntz et al. (2015) - Repeat until no |EE| >= mu*_tresh
if np.all(mustar < mu_thresh): donext = False
niter += 1
# Step 7 of Cuntz et al. (2015) - last screening with ntlast trajectories
# all parameters with mu* < mu*_thresh are final noninformative parameters
if crank == 0:
if (verbose > 0):
tee('\nParameters remaining for last screening:', np.where(imask)[0] + 1, file=lfile)
else:
if lfile is not None:
print('\nParameters remaining for last screening:', np.where(imask)[0] + 1, file=lfile)
iimask = np.where(imask)[0]
res = screening( # (npara,3) with mu*, mu, std if nt>1
func, lb, ub, ntlast,
x0=ix0, mask=imask,
nsteps=nsteps, ntotal=10 * ntlast,
processes=processes, pool=ipool,
verbose=0)
if res.ndim > 2:
if weight:
mustar = np.sum(res[:, iimask, 2] * res[:, iimask, 0], axis=0) / np.sum(res[:, iimask, 2], axis=0)
else:
mustar = np.mean(res[:, iimask, 0], axis=0)
else:
mustar = res[iimask, 0]
if crank == 0:
if ntlast > 1:
if (verbose > 0):
tee('Final mu*: ', mustar, file=lfile)
else:
if lfile is not None:
print('Final mu*: ', mustar, file=lfile)
else:
if (verbose > 0):
tee('Final absolute elementary effects |EE|: ', mustar, file=lfile)
else:
if lfile is not None:
print('Final absolute elementary effects |EE|: ', mustar, file=lfile)
imask[iimask] = imask[iimask] & (mustar < mu_thresh)
if np.all(~imask):
if crank == 0:
if (verbose > 0):
tee('\nNo more parameters left after screening, i.e. all (unmasked) parameters are informative.',
file=lfile)
tee('Finished screening in eee.', file=lfile)
else:
if lfile is not None:
print('\nNo more parameters left after screening, i.e. all (unmasked) parameters are informative.',
file=lfile)
print('Finished screening in eee.', file=lfile)
_cleanup(lfile, pool, ipool)
# Return all true
if mask is None:
return np.ones(len(lb), dtype=np.bool)
else:
return mask
# Return mask with unmasked informative model parameters (to be used with 'and' on initial mask)
if mask is None:
out = ~imask
else:
out = (~imask) & mask # (true where now zero, i.e. were masked or informative) and (initial mask)
if crank == 0:
if (verbose > 0):
tee('\nFinal informative parameters:', np.sum(out), np.where(out)[0] + 1, file=lfile)
tee('Final noninformative parameters:', np.sum(imask), np.where(imask)[0] + 1, file=lfile)
tee('\nFinished screening in eee.', file=lfile)
else:
if lfile is not None:
print('\nFinal informative parameters:', np.sum(out), np.where(out)[0] + 1, file=lfile)
print('Final noninformative parameters:', np.sum(imask), np.where(imask)[0] + 1, file=lfile)
print('\nFinished screening in eee.', file=lfile)
# Close logfile and pool
_cleanup(lfile, pool, ipool)
return out
# across both training points and Hamiltonians.
# For serial jobs, this will make no difference.
tasks = [(name, i, p, True) for name in hamiltonians for i, p in enumerate(p_train)]
tasks_validation = [
(name, i, p, False) for name in hamiltonians for i, p in enumerate(p_valid)
]
def run_evc(task):
# A hacky way to get MPI to play nice with a callable object
# Must be defined before creating the pool
# https://github.com/dfm/emcee/issues/199
return ham_trainer(task)
pool = schwimmbad.choose_pool(mpi=use_mpi)
with pool:
if not pool.is_master():
# Do not repeat the actions on each child process.
# Let the main process distribute the jobs.
pool.wait()
sys.exit(0)
# Create the output directory if it doesn't already exist
if not os.path.exists(output_directory):
os.makedirs(output_directory)
# Put the parameters file with the output so we know how it was generated.
shutil.copy2(parameters_file, output_directory)
print("Starting tasks", flush=FLUSH)
else: # default
logger.setLevel(logging.INFO)
joker_logger.setLevel(logging.INFO)
# Load data here
filename = path.join(args.data_path,
'{0}.{1}'.format(args.apogee_id, args.data_file_ext))
data_tbl = QTable.read(filename)
data = RVData(t=data_tbl['time'],
rv=data_tbl['rv'],
stddev=data_tbl['rv_err'])
# parse config file
with open(args.config_file, 'r') as f:
config = yaml.load(f.read())
config['config_file'] = args.config_file
joker_params = config_to_jokerparams(config)
pool_kwargs = dict(mpi=args.mpi, processes=args.n_procs)
pool = choose_pool(**pool_kwargs)
main(data_path=args.data_path,
pool=pool,
apogee_id=args.apogee_id,
config=config,
overwrite=args.overwrite,
data_file_ext=args.data_file_ext)
pool.close()
sys.exit(0)
def cli():
"""Command-line interface
"""
import argparse
# Help string to be shown using the -h option
descStr = """
Produce common resolution cubes for QUOCKA data.
A seperate cube per band will be produced, along with frequency text files.
"""
# Parse the command line options
parser = argparse.ArgumentParser(
description=descStr, formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument(
'datadir',
metavar='datadir',
type=str,
help='Directory containing a single QUOCKA field images.')
parser.add_argument('field',
metavar='field',
type=str,
help='QUOCKA field name.')
parser.add_argument(
'-o',
'--outdir',
dest='outdir',
type=str,
default=None,
help='(Optional) Save cubes to different directory [datadir].')
parser.add_argument(
'-c',
'--cutoff',
dest='cutoff',
type=float,
default=15,
help="Flags channels with BMAJ > cutoff in arcsec [15]")
parser.add_argument("-v",
"--verbose",
dest="verbose",
action="store_true",
help="verbose output [False].")
parser.add_argument("-d",
"--dryrun",
dest="dryrun",
action="store_true",
help="Compute common beam and stop [False].")
parser.add_argument("--debug",
dest="debug",
action="store_true",
help="Show debugging plots [False].")
parser.add_argument("-t",
"--tolerance",
dest="tolerance",
type=float,
default=0.0001,
help="tolerance for radio_beam.commonbeam.")
parser.add_argument("-e",
"--epsilon",
dest="epsilon",
type=float,
default=0.0005,
help="epsilon for radio_beam.commonbeam.")
parser.add_argument("-n",
"--nsamps",
dest="nsamps",
type=int,
default=200,
help="nsamps for radio_beam.commonbeam.")
group = parser.add_mutually_exclusive_group()
group.add_argument("--ncores",
dest="n_cores",
default=1,
type=int,
help="Number of processes (uses multiprocessing).")
group.add_argument("--mpi",
dest="mpi",
default=False,
action="store_true",
help="Run with MPI.")
args = parser.parse_args()
pool = schwimmbad.choose_pool(mpi=args.mpi, processes=args.n_cores)
if args.mpi:
if not pool.is_master():
pool.wait()
sys.exit(0)
# make it so we can use imap in serial and mpi mode
if not isinstance(pool, schwimmbad.MultiPool):
pool.imap = pool.map
verbose = args.verbose
main(pool, args, verbose=verbose)
pool.close()
def cli():
"""Command-line interface
"""
import argparse
# Help string to be shown using the -h option
descStr = """
Smooth a field of 2D images to a common resolution.
Names of output files are 'infile'.sm.fits
"""
# Parse the command line options
parser = argparse.ArgumentParser(
description=descStr, formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument(
'infile',
metavar='infile',
type=str,
help=
'Input FITS image to smooth (can be a wildcard) - beam info must be in header.'
)
parser.add_argument('-p',
'--prefix',
dest='prefix',
type=str,
default=None,
help='Add prefix to output filenames.')
parser.add_argument(
'-o',
'--outdir',
dest='outdir',
type=str,
default=None,
help='Output directory of smoothed FITS image(s) [./].')
parser.add_argument("-v",
"--verbose",
dest="verbose",
action="store_true",
help="verbose output [False].")
parser.add_argument(
"--bmaj",
dest="bmaj",
type=float,
default=None,
help="BMAJ to convolve to [max BMAJ from given image(s)].")
parser.add_argument(
"--bmin",
dest="bmin",
type=float,
default=None,
help="BMIN to convolve to [max BMAJ from given image(s)].")
parser.add_argument("--bpa",
dest="bpa",
type=float,
default=None,
help="BPA to convolve to [0].")
group = parser.add_mutually_exclusive_group()
group.add_argument("--ncores",
dest="n_cores",
default=1,
type=int,
help="Number of processes (uses multiprocessing).")
group.add_argument("--mpi",
dest="mpi",
default=False,
action="store_true",
help="Run with MPI.")
args = parser.parse_args()
pool = schwimmbad.choose_pool(mpi=args.mpi, processes=args.n_cores)
if args.mpi:
if not pool.is_master():
pool.wait()
sys.exit(0)
verbose = args.verbose
main(pool, args, verbose=verbose)
pool.close()
hugs.utils.mkdir_if_needed(outdir)
log_dir = os.path.join(outdir, 'log')
hugs.utils.mkdir_if_needed(log_dir)
log_fn = []
for tract, patch in patches['tract', 'patch']:
fn = os.path.join(log_dir, '{}-{}.log'.format(tract, patch))
log_fn.append(fn)
patches['outdir'] = outdir
patches['log_fn'] = log_fn
else:
print('\n**** must give tract and patch --or-- a patch file ****\n')
parser.print_help()
exit()
patches['rerun_path'] = args.rerun_path
patches['seed'] = args.seed
patches['config_fn'] = args.config_fn
patches['run_name'] = args.run_name
if rank==0:
# open database session with master process
db_fn = os.path.join(outdir, args.run_name+'.db')
engine = hugs.database.connect(db_fn, args.overwrite)
session = hugs.database.Session()
shutil.copyfile(args.config_fn, os.path.join(outdir, 'config.yml'))
pool = schwimmbad.choose_pool(mpi=args.mpi, processes=args.ncores)
list(pool.map(worker, patches, callback=ingest_data))
pool.close()
def screening(func,
lb,
ub,
nt,
x0=None,
mask=None,
nsteps=6,
ntotal=None,
seed=None,
processes=1,
pool=None,
verbose=0):
"""
Parameter screening using Morris method of Elementary Effects.
Note, the input function must be callable as `func(x)`.
Parameters
----------
func : callable
Python function callable as `func(x)` with `x` the function parameters.
lb : array_like
Lower bounds of parameters.
ub : array_like
Upper bounds of parameters.
nt : int
Number of trajectories used for screening.
x0 : array_like, optional
Parameter values used with `mask==0`.
mask : array_like, optional
Include (1,True) or exclude (0,False) parameters in screening (default: include all parameters).
nsteps : int, optional
Number of steps along one trajectory (default: 6)
ntotal : int, optional
Total number of sampled trajectories to select the nt most different trajectories.
If None: `max(nt**2,10*nt)` (default: None)
seed : int or array_like
Seed for numpy``s random number generator (default: None).
processes : int, optional
The number of processes to use to evaluate objective function and constraints (default: 1).
pool : `schwimmbad` pool object, optional
Generic map function used from module `schwimmbad <https://schwimmbad.readthedocs.io/en/latest/>`_,
which provides, serial, multiprocessor, and MPI mapping functions (default: None).
The pool is chosen with:
schwimmbad.choose_pool(mpi=True/False, processes=processes).
The pool will be chosen automatically if pool is None.
MPI pools can only be opened and closed once. If you want to use screening several
times in one program, then you have to choose the pool, pass it to screening,
and later close the pool in the calling progam.
verbose : int, optional
Print progress report during execution if verbose>0 (default: 0).
Returns
-------
(nparameter,3) ndarray
if nt>1:
2D-array - (nparameter,3) with per parameter
1. mean of absolute elementary effects over all nt trajectories (mu*)
2. mean of elementary effects over all nt trajectories (mu)
3. standard deviation of elementary effects over all nt trajectories (sigma)
else:
2D-array - (nparameter,3) with per parameter
1. absolute elementary effect of each parameter
2. elementary effect of each parameter
3. zeros
See Also
--------
:func:`~pyeee.eee.eee` : Efficient Elementary Effects, same as
:func:`~pyeee.eee.see` : Sequential Elementary Effects
Examples
--------
>>> from functools import partial
>>> import numpy as np
>>> from function_wrapper import func_wrapper
>>> from sa_test_functions import fmorris
>>> seed = 1023
>>> np.random.seed(seed=seed)
>>> npars = 20
>>> beta0 = 0.
>>> beta1 = np.random.standard_normal(npars)
>>> beta1[:10] = 20.
>>> beta2 = np.random.standard_normal((npars,npars))
>>> beta2[:6,:6] = -15.
>>> beta3 = np.zeros((npars,npars,npars))
>>> beta3[:5,:5,:5] = -10.
>>> beta4 = np.zeros((npars,npars,npars,npars))
>>> beta4[:4,:4,:4,:4] = 5.
>>> arg = [beta0, beta1, beta2, beta3, beta4]
>>> kwarg = {}
>>> func = partial(func_wrapper, fmorris, arg, kwarg)
>>> lb = np.zeros(npars)
>>> ub = np.ones(npars)
>>> nt = 20
>>> ntotal = 10*nt
>>> nsteps = 6
>>> verbose = 0
>>> out = screening(func, lb, ub, nt, x0=None, mask=None, nsteps=nsteps, ntotal=ntotal, processes=4, verbose=verbose)
>>> print(out[0:3,0])
[60.7012889 67.33372626 48.46673528]
History
-------
Written, Matthias Cuntz, Dec 2017
Modified, Matthias Cuntz, Dec 2017 - output for nt=1 also (npara,3)
Matthias Cuntz, Jan 2018 - rm call of external programs
Matthias Cuntz, Jan 2018 - function can return multiple output, e.g. time series
Fabio Gennaretti, Jul 2018 - map of Python3 returns iterator -> make list(map())
Matthias Cuntz, Dec 2019 - bug: default ntotal was not set if ntotal<0 (but nt instead)
- numpy docstring format
Matthias Cuntz, Jan 2020 - x0 optional
- distinguish iterable and array_like parameter types
- added seed
Matthias Cuntz, Feb 2020 - InputError -> TypeError
- use new names of kwargs of moris_sampling
Matthias Cuntz, Feb 2020 - number of final trajectories is argument, not keyword argument
"""
# Get MPI communicator
try:
from mpi4py import MPI
comm = MPI.COMM_WORLD
csize = comm.Get_size()
crank = comm.Get_rank()
except ImportError:
comm = None
csize = 1
crank = 0
# Checks
# Bounds
assert len(lb) == len(
ub), 'Lower- and upper-bounds must have the same lengths.'
lb = np.array(lb)
ub = np.array(ub)
# Mask
if mask is None:
assert np.all(
ub >= lb
), 'All upper-bounds must be greater or equal than lower-bounds.'
else:
assert len(mask) == len(
ub), 'Mask and bounds must have the same lengths.'
if x0 is None:
raise TypeError('x0 must be given if mask is set')
x0 = np.array(x0)
if not np.all(mask):
assert len(mask) == len(
x0), 'Mask and x0 must have the same lengths.'
assert np.all(
ub[mask] >= lb[mask]
), 'All unmasked upper-bounds must be greater or equal than lower-bounds.'
# Set defaults
npara = len(lb)
if mask is None:
imask = np.ones(npara, dtype=np.bool)
else:
imask = mask
nmask = np.sum(imask)
if ntotal is None:
ntotal = max(nt**2, 10 * nt)
# Seed random number generator
if seed is not None:
np.random.seed(
seed=seed
) # same on all ranks because trajectories are sampled on all ranks
# Sample trajectories
if (crank == 0) and (verbose > 0): print('Sample trajectories')
tmatrix, tvec = morris_sampling(nmask,
lb[imask],
ub[imask],
nt,
nsteps=nsteps,
ntotal=ntotal,
Diagnostic=False)
if mask is None:
x = tmatrix
else:
# Set input vector to trajectories and masked elements = x0
x = np.tile(x0, tvec.size).reshape(tvec.size, npara) # default to x0
x[:, imask] = tmatrix # replaced unmasked with trajectorie values
# Choose the right mapping function: single, multiprocessor or mpi
if pool is None:
import schwimmbad
ipool = schwimmbad.choose_pool(mpi=False if csize == 1 else True,
processes=processes)
else:
ipool = pool
# Calculate all model runs
if (crank == 0) and (verbose > 0): print('Calculate objective functions')
fx = np.array(list(ipool.map(func, x)))
# Calc elementary effects
if (crank == 0) and (verbose > 0): print('Calculate elementary effects')
if fx.ndim == 1:
fx = fx[:, np.newaxis]
nfx = 1
else:
nfx = fx.shape[1]
out = np.zeros((nfx, npara, 3))
for j in range(nfx):
sa, res = elementary_effects(nmask,
tmatrix,
tvec,
fx[:, j],
nsteps=nsteps,
Diagnostic=False)
# Output with zero for all masked parameters
if nt == 1:
out[j, imask, 0] = np.abs(sa[:, 0])
out[j, imask, 1] = sa[:, 0]
else:
out[j, imask, :] = res
if nfx == 1: out = out[0, :, :]
if pool is None: ipool.close()
return out
def log_prob_fn(p):
# A hacky way to get MPI to play nice with a callable object
# Must be defined before creating the pool
# https://github.com/dfm/emcee/issues/199
return obs_dist(p)
# We must define the pool *after* the log_prob_fn
# Otherwise MPI will complain
# If the --mpi flag is passed, then this will
# return an MPIPool object, otherwise it will
# be a SerialPool object. This helps us test the
# script before going parallel.
pool = schwimmbad.choose_pool(mpi=args.mpi)
####################################################################
# Set up the sampler and run it!
####################################################################
with pool:
if not pool.is_master():
# Do not repeat the actions on each child process.
# Let the main process distribute the jobs.
pool.wait()
sys.exit(0)
# Create the output directory if it doesn't already exist
if not os.path.exists(output_directory):
os.makedirs(output_directory)
print('running metadetect', flush=True)
print('config:', CONFIG, flush=True)
print('swap 12:', SWAP12)
print('use mpi:', USE_MPI, flush=True)
print("n_ranks:", n_ranks, flush=True)
print("n_workers:", n_workers, flush=True)
if n_workers == 1:
outputs = [_run_sim(0)]
else:
if not USE_MPI:
pool = schwimmbad.JoblibPool(n_workers,
backend='multiprocessing',
verbose=100)
else:
pool = schwimmbad.choose_pool(mpi=USE_MPI, processes=n_workers)
outputs = pool.map(_run_sim, range(n_sims))
pool.close()
pres, mres = zip(*outputs)
pres, mres = cut_nones(pres, mres)
if rank == 0:
dt = [('g1p', 'f8'), ('g1m', 'f8'), ('g1', 'f8'), ('g2p', 'f8'),
('g2m', 'f8'), ('g2', 'f8')]
dplus = np.array(pres, dtype=dt)
dminus = np.array(mres, dtype=dt)
with fitsio.FITS('data.fits', 'rw') as fits:
fits.write(dplus, extname='plus')
fits.write(dminus, extname='minus')
def _generate_koi_maxAV_table(procs=1):
kois = np.array(ku.DR25.index)
pool = choose_pool(mpi=False, processes=procs)
maxAV = pool.map(_getAV, kois)
np.savetxt(KOI_MAXAV_FILE, np.array([kois, maxAV]).T, fmt='%s %.3f')
def swim(self, lens_name, model_id, log=True, mpi=False,
recipe_name='default', sampler='EMCEE', thread_count=1):
"""
Run models for a single lens.
:param lens_name: lens name
:type lens_name: `str`
:param model_id: identifier for the model run
:type model_id: `str`
:param log: if `True`, all `print` statements will be logged
:type log: `bool`
:param mpi: MPI option
:type mpi: `bool`
:param recipe_name: recipe for pre-sampling optimization, supported
ones now: 'default' and 'galaxy-galaxy'
:type recipe_name: `str`
:param sampler: 'EMCEE' or 'COSMOHAMMER', cosmohammer is kept for
legacy
:type sampler: `str`
:param thread_count: number of threads if `multiprocess` is used
:type thread_count: `int`
:return:
:rtype:
"""
pool = choose_pool(mpi=mpi)
if log and pool.is_master():
log_file = open(self.file_system.get_log_file_path(lens_name,
model_id), 'wt')
sys.stdout = log_file
config = self.get_lens_config(lens_name)
recipe = Recipe(config, sampler=sampler, thread_count=thread_count)
psf_supersampling_factor = config.get_psf_supersampled_factor()
kwargs_data_joint = self.get_kwargs_data_joint(
lens_name,
psf_supersampled_factor=psf_supersampling_factor)
fitting_sequence = FittingSequence(
kwargs_data_joint,
config.get_kwargs_model(),
config.get_kwargs_constraints(),
config.get_kwargs_likelihood(),
config.get_kwargs_params(),
mpi=mpi
)
fitting_kwargs_list = recipe.get_recipe(
kwargs_data_joint=kwargs_data_joint,
recipe_name=recipe_name)
fit_output = fitting_sequence.fit_sequence(fitting_kwargs_list)
kwargs_result = fitting_sequence.best_fit(bijective=False)
output = {
'settings': config.settings,
'kwargs_result': kwargs_result,
'fit_output': fit_output,
}
if pool.is_master():
self.file_system.save_output(lens_name, model_id, output)
if log and pool.is_master():
log_file.close()
def main():
import argparse
"""
Start the function to perform RM-clean if called from the command line.
"""
# Help string to be shown using the -h option
descStr = """
Run RM-CLEAN on a cube of Faraday dispersion functions (FDFs), applying
a cube of rotation measure spread functions created by the script
'do_RMsynth_3D.py'. Saves cubes of deconvolved FDFs & clean-component
spectra, and a pixel map showing the number of iterations performed.
Set any of the multiprocessing options to enable parallelization
(otherwise, pixels will be processed serially).
Expects that the input is in the form of the Stokes-separated
(single extension) FITS cubes produced by do_RMsynth_3D.
"""
# Parse the command line options
parser = argparse.ArgumentParser(
description=descStr, formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument(
"fitsFDF",
metavar="FDF_dirty.fits",
nargs=1,
help=
"FITS cube containing the dirty FDF.\n(Can be any of the FDF output cubes from do_RMsynth_3D.py)"
)
parser.add_argument(
"fitsRMSF",
metavar="RMSF.fits",
nargs=1,
help=
"FITS cube containing the RMSF and FWHM image.\n(Cans be any of the RMSF output cubes from do_RMsynth_3D.py)"
)
parser.add_argument("-c",
dest="cutoff",
type=float,
nargs=1,
default=1.0,
help="CLEAN cutoff in flux units")
parser.add_argument(
"-n",
dest="maxIter",
type=int,
default=1000,
help="Maximum number of CLEAN iterations per pixel [1000].")
parser.add_argument("-g",
dest="gain",
type=float,
default=0.1,
help="CLEAN loop gain [0.1].")
parser.add_argument("-o",
dest="prefixOut",
default="",
help="Prefix to prepend to output files [None].")
parser.add_argument(
"-f",
dest="write_separate_FDF",
action="store_false",
help=
"Store different Stokes as FITS extensions [False, store as separate files]."
)
parser.add_argument("-v",
dest="verbose",
action="store_true",
help="Verbose [False].")
group = parser.add_mutually_exclusive_group()
group.add_argument("--ncores",
dest="n_cores",
default=1,
type=int,
help="Number of processes (uses multiprocessing).")
parser.add_argument(
"--chunk",
dest="chunk",
default=None,
type=int,
help="Chunk size (uses multiprocessing -- not available in MPI!)")
group.add_argument("--mpi",
dest="mpi",
default=False,
action="store_true",
help="Run with MPI.")
args = parser.parse_args()
#Check if the user is trying to use parallelization without installing it:
if parallel_available == False and (args.n_cores != 1 or args.chunk != None
or args.mpi != False):
raise Exception(
'Parallel processing not available. Please install the schwimmbad module to enable parallel processing.'
)
#If parallelization requested use it, otherwise use the old-fashioned way.
if parallel_available == True and (args.n_cores != 1 or args.chunk != None
or args.mpi != False):
pool = schwimmbad.choose_pool(mpi=args.mpi, processes=args.n_cores)
if args.mpi:
if not pool.is_master():
pool.wait()
sys.exit(0)
if args.n_cores > 1:
chunksize = args.chunk
else:
chunksize = None
else:
pool = None
chunksize = None
verbose = args.verbose
# Sanity checks
for f in args.fitsFDF + args.fitsRMSF:
if not os.path.exists(f):
print("File does not exist: '%s'." % f)
sys.exit()
dataDir, dummy = os.path.split(args.fitsFDF[0])
# Run RM-CLEAN on the cubes
cleanFDF, ccArr, iterCountArr, residFDF, headtemp = run_rmclean(
fitsFDF=args.fitsFDF[0],
fitsRMSF=args.fitsRMSF[0],
cutoff=args.cutoff,
maxIter=args.maxIter,
gain=args.gain,
chunksize=chunksize,
nBits=32,
verbose=verbose)
# Write results to disk
writefits(cleanFDF,
ccArr,
iterCountArr,
residFDF,
headtemp,
prefixOut=args.prefixOut,
outDir=dataDir,
write_separate_FDF=args.write_separate_FDF,
verbose=verbose)
for r in pool.map(worker, [row for _, row in rgbs.iterrows()],
callback=worker.write_results):
pass
pool.close()
if __name__ == '__main__':
from argparse import ArgumentParser
from isochrones import get_ichrone
parser = ArgumentParser(description="Run stellar model fits.")
parser.add_argument('--nrows', dest='n_rows', type=int, default=None)
group = parser.add_mutually_exclusive_group()
group.add_argument("--ncores",
dest="n_cores",
default=1,
type=int,
help="Number of processes (uses multiprocessing).")
group.add_argument("--mpi",
dest="mpi",
default=False,
action="store_true",
help="Run with MPI.")
args = parser.parse_args()
pool = schwimmbad.choose_pool(mpi=args.mpi, processes=args.n_cores)
main(pool, n_rows=args.n_rows)
def main(pool_kwargs):
# Make some fake data
orbit = KeplerOrbit(P=30 * u.day, e=0.1, M0=0 * u.rad, omega=0 * u.deg)
K = 1 * u.km / u.s
samples0 = JokerSamples()
samples0['P'] = 30 * u.day
samples0['e'] = 0.1 * u.one
samples0['M0'] = 0 * u.rad
samples0['omega'] = 0 * u.rad
samples0['jitter'] = 1 * u.m / u.s
samples0['K'] = 1 * u.km / u.s
samples0['v0'] = 0 * u.km / u.s
t = Time(59000 + np.linspace(0, 100, 16), format='mjd')
err = np.ones_like(t.mjd) * 10 * u.m / u.s
data = RVData(t=t, rv=K * orbit.unscaled_radial_velocity(t), stddev=err)
params = JokerParams(P_min=1 * u.day, P_max=1024 * u.day)
pool = choose_pool(**pool_kwargs)
model = TheJokerMCMCModel(joker_params=params, data=data)
p0_mean = np.squeeze(model.pack_samples(samples0))
ball_std = 1E-3 # TODO: make this customizable?
# P, M0, e, omega, jitter, K, v0
n_walkers = 1024
p0 = np.zeros((n_walkers, len(p0_mean)))
for i in range(p0.shape[1]):
if i in [2, 4]: # eccentricity, jitter
p0[:, i] = np.abs(
np.random.normal(p0_mean[i], ball_std, size=n_walkers))
else:
p0[:, i] = np.random.normal(p0_mean[i], ball_std, size=n_walkers)
p0 = model.to_mcmc_params(p0.T).T
# Because jitter is always carried through in the transform above, now
# we have to remove the jitter parameter if it's fixed!
if params._fixed_jitter:
p0 = np.delete(p0, 5, axis=1)
# time0 = time.time()
# results = list(pool.map(model,
# # results = list(pool.map(test_ln_prob,
# (p0[i] for i in range(len(p0)))))
# print('...time spent mapping: {0}'.format(time.time()-time0))
#
# pool.close()
# sys.exit(0)
n_dim = p0.shape[1]
sampler = emcee.EnsembleSampler(n_walkers,
n_dim,
make_model_ln_prob,
pool=pool,
args=(params, data))
n_burn = 128
print('Burning in MCMC for {0} steps...'.format(n_burn))
time0 = time.time()
pos, *_ = sampler.run_mcmc(p0, n_burn)
print('...time spent burn-in: {0}'.format(time.time() - time0))
pool.close()
sys.exit(0)
parser.add_argument("-o",
"--overwrite",
dest="overwrite",
default=False,
action="store_true",
help="Overwrite stuff.")
parser.add_argument("-n",
"--name",
dest="name",
required=True,
choices=['ms', 'rg', 'ms_mh', 'rg_mh'])
args = parser.parse_args()
pool = choose_pool(mpi=args.mpi)
# Make sure paths exist:
for path_ in [cache_path, plot_path]:
os.makedirs(path_, exist_ok=True)
func = eval('main_{}'.format(args.name))
try:
func(pool=pool, overwrite=args.overwrite)
pool.close()
except Exception as e:
pool.close()
raise e
sys.exit(0)
def main(n_sims, seed, output_file, serial):
"""Run N_SIMS metadetect simulation patches and estimate the shear."""
# logging and MPI/workers
_deal_with_logging(n_sims)
use_mpi, comm, rank, n_ranks, n_workers = _deal_with_mpi(n_sims)
# and the config
(sim_config, run_config, shear_meas_config, swap12,
cut_interp) = load_config('config.yaml')
use_old_sim = sim_config.pop('straight_to_coadd', False)
if use_old_sim:
sim_class = SimpleSim
else:
sim_class = CoaddingSim
__run_sim = partial(_run_sim,
sim_config=sim_config,
shear_meas_config=shear_meas_config,
swap12=swap12,
cut_interp=cut_interp,
sim_class=sim_class)
if rank == 0:
print('running metadetect', flush=True)
print('config:', sim_config, flush=True)
print('swap 12:', swap12)
print('use mpi:', use_mpi, flush=True)
print("n_ranks:", n_ranks, flush=True)
print("n_workers:", n_workers, flush=True)
if use_old_sim:
print('sim type: straight-to-coadd')
assert False, (
"Straight-to-coadd sims are disbaled for safety! Remove this "
"assert if you actually want to use it!")
else:
print('sim type: full coadding')
seeds = np.random.RandomState(seed=seed).randint(low=1,
high=2**30,
size=n_sims).tolist()
if n_workers == 1:
outputs = [__run_sim(seeds[0])]
elif serial:
outputs = [__run_sim(seed) for seed in seeds]
else:
if not use_mpi:
pool = schwimmbad.JoblibPool(n_workers,
backend='multiprocessing',
verbose=40)
else:
pool = schwimmbad.choose_pool(mpi=use_mpi, processes=n_workers)
outputs = pool.map(__run_sim, seeds)
pool.close()
outputs = [o for o in outputs if o is not None]
pres, mres = zip(*outputs)
all_data = {}
for s2n in S2N_CUTS:
data_p = np.concatenate([p[s2n] for p in pres])
data_m = np.concatenate([m[s2n] for m in mres])
all_data['p%d' % s2n] = data_p
all_data['m%d' % s2n] = data_m
m, msd, c, csd = estimate_m_and_c(data_p,
data_m,
0.02,
swap12=swap12,
step=0.01,
rng=np.random.RandomState(seed=42),
n_boot=500,
verbose=False)
print("""\
s2n: {s2n}
# of sims: {n_sims}
noise cancel m : {m:f} +/- {msd:f}
noise cancel c : {c:f} +/- {csd:f}""".format(n_sims=len(data_p),
m=m,
msd=msd,
c=c,
csd=csd,
s2n=s2n),
flush=True)
with fitsio.FITS(output_file, 'rw', clobber=True) as fits:
for s2n in S2N_CUTS:
for s in ['p', 'm']:
ext = '%s%d' % (s, s2n)
fits.write(all_data[ext], extname=ext)
