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)