text:
import numpy as np

align_mat = np.array([
    [-1.96305411e-04,  6.56927099e-06, 9.99999981e-01],
    [ 9.99994245e-01, -3.38686178e-03, 1.96326534e-04],
    [ 3.38686300e-03, 9.99994265e-01, -5.90437390e-06]])

def maintain_quat(q, round_off_err=1e-6):

    norm = np.sqrt(np.sum(np.square(q)))
    if (q[3] < 0) or (abs(norm - 1.0) > round_off_err):
        # print "fixing quat"
        invnorm = 1./norm
        err = abs(norm - 1.0)
        if q[3] < 0:
            invnorm = -invnorm
        q = q*invnorm
    return q

def convertRotMatrixToQuat(rot):

    diag_sum = np.zeros(4)
    q = np.zeros(4)

    diag_sum[0]=1+rot[0][0]-rot[1][1]-rot[2][2]
    diag_sum[1]=1-rot[0][0]+rot[1][1]-rot[2][2]
    diag_sum[2]=1-rot[0][0]-rot[1][1]+rot[2][2]
    diag_sum[3]=1+rot[0][0]+rot[1][1]+rot[2][2]

    maxi = np.argmax(diag_sum)

    q[maxi] = .5*np.sqrt(diag_sum[maxi])
    recip = 1./(4.*q[maxi])

    if maxi == 0:
        q[1] = recip*(rot[0,1] + rot[1,0])
        q[2] = recip*(rot[2,0] + rot[0,2])
        q[3] = recip*(rot[1,2] - rot[2,1])

    elif maxi == 1:
        q[0] = recip*(rot[0,1] + rot[1,0])
        q[2] = recip*(rot[1,2] + rot[2,1])
        q[3] = recip*(rot[2,0] - rot[0,2])

    elif maxi == 2:
        q[0] = recip*(rot[2,0] + rot[0,2])
        q[1] = recip*(rot[1,2] + rot[2,1])
        q[3] = recip*(rot[0,1] - rot[1,0])

    elif maxi == 3:
        q[0] = recip*(rot[1,2] - rot[2,1])
        q[1] = recip*(rot[2,0] - rot[0,2])
        q[2] = recip*(rot[0,1] - rot[1,0])

    q = maintain_quat(q)

    return q

def convertQuatToRotMatrix(q):

    rot = np.zeros((3,3))
    q2 = np.square(q)

    rot[0,0] = q2[0] - q2[1] - q2[2] + q2[3]
    rot[1,1] = -q2[0] + q2[1] - q2[2] + q2[3]
    rot[2,2] = -q2[0] - q2[1] + q2[2] + q2[3]

    rot[0,1] = 2.*(q[0]*q[1] + q[2]*q[3])
    rot[1,0] = 2.*(q[0]*q[1] - q[2]*q[3])

    rot[0,2] = 2.*(q[0]*q[2] - q[1]*q[3])
    rot[2,0] = 2.*(q[0]*q[2] + q[1]*q[3])

    rot[1,2] = 2.*(q[1]*q[2] + q[0]*q[3])
    rot[2,1] = 2.*(q[1]*q[2] - q[0]*q[3])

    return rot

def invertQuat(q):

    q1 = -1.*q
    q1[3] *= -1

    return q1

def productOfQuats(q1, q2):

    q = np.zeros_like(q1)

    q[0] =  q2[3]*q1[0] + q2[2]*q1[1] - q2[1]*q1[2] + q2[0]*q1[3]
    q[1] = -q2[2]*q1[0] + q2[3]*q1[1] + q2[0]*q1[2] + q2[1]*q1[3]
    q[2] =  q2[1]*q1[0] - q2[0]*q1[1] + q2[3]*q1[2] + q2[2]*q1[3]
    q[3] = -q2[0]*q1[0] - q2[1]*q1[1] - q2[2]*q1[2] + q2[3]*q1[3]

    q = maintain_quat(q)

    return q

def applyRotMatrixToVector(rot, vec):

    newvec = np.zeros_like(vec)

    for i in range(len(newvec)):
        if vec.ndim > 1:
            newvec[i] = np.sum(vec*rot[i,:,np.newaxis], axis=0)
        else:
            newvec[i] = np.sum(vec*rot[i,:])

    return newvec

def radec2skyvec(ra, dec):
    z_sky = np.sin(np.radians(dec))
    x_sky = (np.cos(np.radians(dec)))*(np.cos(np.radians(ra)))
    y_sky = (np.cos(np.radians(dec)))*(np.sin(np.radians(ra)))
    skyunit = np.array([x_sky,y_sky,z_sky])
    return skyunit

def skyvec2radec(sky_vec):
    dec = np.rad2deg(np.arcsin(sky_vec[2]))
    ra = np.rad2deg(np.arctan2(sky_vec[1],sky_vec[0]))
    if np.isscalar(ra):
        if ra < 0:
            ra += 360.0
    else:
        if np.any(ra<0):
            ra[(ra<0)] += 360.0
    return ra, dec

def imxy2detvec(imx, imy):

    z= 1./np.sqrt(imx**2. + imy**2. + 1.)
    x = imx*z
    y= imy*z
    detunit = np.array([x,y,z])
    return detunit

def theta_phi2detvec(theta, phi):

    z = np.cos(np.radians(theta))
    x = np.sin(np.radians(theta))*np.cos(np.radians(phi))
    y = np.sin(np.radians(theta))*np.sin(np.radians(-phi))
    detunit = np.array([x,y,z])
    return detunit

def convertSkyToSensor(teldef_qinv, sky_vec, att_q):

    cor_q = productOfQuats(att_q, teldef_qinv)

    rot = convertQuatToRotMatrix(cor_q)

    det_vec = applyRotMatrixToVector(rot, sky_vec)

    return det_vec

def convertDetToSky(teldef_qinv, det_vec, att_q):

    cor_q = productOfQuats(att_q, teldef_qinv)

    cor_qinv = invertQuat(cor_q)

    rot = convertQuatToRotMatrix(cor_qinv)

    sky_vec = applyRotMatrixToVector(rot, det_vec)

    return sky_vec

def convert_radec2batxyz(ra, dec, att_q):
    """
    Converts ra, dec to the BAT x,y,z vector

    Parameters:
    ra: float in degrees
    dec: float in degrees
    att_q: the attitude quaternion (length 4 1D array)

    Returns:
    detector vec: length 3 float vector
    """
    teldef_q = convertRotMatrixToQuat(align_mat)
    teldef_qinv = invertQuat(teldef_q)
    att_q = maintain_quat(att_q)

    sky_vec = radec2skyvec(ra, dec)
    det_vec = convertSkyToSensor(teldef_qinv, sky_vec, att_q)

    return det_vec

def convert_radec2thetaphi(ra, dec, att_q):
    """
    Converts ra, dec to the BAT theta, phi

    Parameters:
    ra: float in degrees
    dec: float in degrees
    att_q: the attitude quaternion (length 4 1D array)

    Returns:
    theta, phi: float in degrees
    """
    teldef_q = convertRotMatrixToQuat(align_mat)
    teldef_qinv = invertQuat(teldef_q)
    att_q = maintain_quat(att_q)

    sky_vec = radec2skyvec(ra, dec)
    det_vec = convertSkyToSensor(teldef_qinv, sky_vec, att_q)

    theta = np.rad2deg(np.arccos(det_vec[2]))
    phi = np.rad2deg(np.arctan2(-det_vec[1], det_vec[0]))
    if np.isscalar(phi):
        if phi < 0.0:
            phi += 360.0
    else:
        bl = (phi<0.0)
        if np.sum(bl) > 0:
            phi[bl] += 360.0

    return theta, phi

def convert_radec2imxy(ra, dec, att_q):
    """
    Converts ra, dec to imx, imy

    Parameters:
    ra: float in degrees
    dec: float in degrees
    att_q: the attitude quaternion (length 4 1D array)

    Returns:
    imx: float
    imy: float
    """
    teldef_q = convertRotMatrixToQuat(align_mat)
    teldef_qinv = invertQuat(teldef_q)
    att_q = maintain_quat(att_q)

    sky_vec = radec2skyvec(ra, dec)
    det_vec = convertSkyToSensor(teldef_qinv, sky_vec, att_q)

    imx = det_vec[0]/det_vec[2]
    imy = det_vec[1]/det_vec[2]
    return imx,imy

def convert_theta_phi2radec(theta, phi, att_q):
    """
    Converts theta, phi to ra, dec

    Parameters:
    theta: float in degrees
    phi: float in degrees
    att_q: the attitude quaternion (length 4 1D array)

    Returns:
    ra: float in degrees
    dec: float in degrees
    """

    teldef_q = convertRotMatrixToQuat(align_mat)
    teldef_qinv = invertQuat(teldef_q)
    att_q = maintain_quat(att_q)

    # det_vec = imxy2detvec(imx, imy)
    det_vec = theta_phi2detvec(theta, phi)

    sky_vec = convertDetToSky(teldef_qinv, det_vec, att_q)

    ra, dec = skyvec2radec(sky_vec)

    return ra, dec

def convert_imxy2radec(imx, imy, att_q):
    """
    Converts imx, imy to ra, dec

    Parameters:
    imx: float
    imy: float
    att_q: the attitude quaternion (length 4 1D array)

    Returns:
    ra: float in degrees
    dec: float in degrees
    """

    teldef_q = convertRotMatrixToQuat(align_mat)
    teldef_qinv = invertQuat(teldef_q)
    att_q = maintain_quat(att_q)

    det_vec = imxy2detvec(imx, imy)

    sky_vec = convertDetToSky(teldef_qinv, det_vec, att_q)

    ra, dec = skyvec2radec(sky_vec)

    return ra, dec

def theta_phi2imxy(theta, phi):

    imr = np.tan(np.radians(theta))
    imx = imr*np.cos(np.radians(phi))
    imy = imr*np.sin(np.radians(-phi))
    return imx, imy

def imxy2theta_phi(imx, imy):

    theta = np.rad2deg(np.arctan(np.sqrt(imx**2 + imy**2)))
    phi = np.rad2deg( np.arctan2(-imy,imx) )
    if np.isscalar(phi):
        if phi < 0:
            phi += 360.0
    else:
        bl = (phi<0)
        if np.sum(bl)>0:
            phi[bl] += 360.0
    return theta, phi

def pnt2euler(ra, dec, roll, roll_sign=-1.0, roll_offset=0.):

    phi = np.radians(ra)
    theta = np.radians(90.-dec)
    psi = np.radians(roll*roll_sign + 90. - roll_offset)

    return phi, theta, psi

def convertEulerToRotMatrix(phi, theta, psi):

    sphi = np.sin(phi); cphi = np.cos(phi)
    stheta = np.sin(theta); ctheta = np.cos(theta)
    spsi = np.sin(psi); cpsi = np.cos(psi)

    rot = np.zeros((3,3))

    rot[0,0] = cpsi*ctheta*cphi - spsi*sphi
    rot[0,1] = cpsi*ctheta*sphi + spsi*cphi

    rot[1,0] = -spsi*ctheta*cphi - cpsi*sphi
    rot[1,1] = -spsi*ctheta*sphi + cpsi*cphi

    rot[0,2] = -cpsi*stheta
    rot[1,2] = spsi*stheta

    rot[2,0] = stheta*cphi
    rot[2,1] = stheta*sphi

    rot[2,2] = ctheta

    return rot

def convertEulerToQuat(phi, theta, psi):

    rot = convertEulerToRotMatrix(phi, theta, psi)

    q = convertRotMatrixToQuat(rot)

    return q

def pnt2quat(ra, dec, roll):
    """
    Converts pointing direciton to a quaternion

    Parameters:
    ra: float in degrees
    dec: float in degrees
    roll: float in degrees

    Returns:
    q: the attitude quaternion (length 4 1D array)
    """

    teldef_q = convertRotMatrixToQuat(align_mat)

    phi, theta, psi = pnt2euler(ra, dec, roll)

    q_realigned = convertEulerToQuat(phi, theta, psi)

    q = productOfQuats(q_realigned, teldef_q)

    return q

max_stars_repo_name: parsotat/NITRATES

max_issues_repo_name: parsotat/NITRATES

max_forks_repo_name: parsotat/NITRATES


text:
""" file: blink.py (maxr.flow)
    author: Jess Robertson
            CSIRO Mineral Resources
    date:   September 25, 2015

    description: Blinking vortex flow
"""

from __future__ import print_function, division

from numpy import asarray, clip, pi, sin, newaxis

def vortex(x_center, y_center, gamma):
    """ Returns a function to give a vortex flow centered at r,
        with strength Gamma
    """
    def _flow(xxs, yys):
        "Wrapped flow function"
        xxs = asarray(xxs)
        yys = asarray(yys)
        rsq = (xxs - x_center) ** 2 + (yys - y_center) ** 2
        return (-gamma * (xxs - x_center) / (2 * pi * rsq),
                gamma * (yys - y_center) / (2 * pi * rsq))
    return _flow

def tick(time, period):
    """ Adding some time dependence!
    """
    return clip(sin(2 * pi * time / period), 0, 1.01)

def tock(time, period):
    """ Adding some time dependence!
    """
    return clip(-sin(2 * pi * time / period), 0, 1.01)

def blink(gamma, period):
    """ The blinking vortex flow
    """
    # Generate two vortices
    vortices = [vortex(x, 0, gamma) for x in (-0.5, 0.5)]

    # Wrap in flow with time dependence
    def _flow(xxs, yys, time):
        "Wrapped flow function"
        _vortex = [v(xxs, yys) for v in vortices]
        return _vortex[0][0][..., newaxis] * tick(time, period) \
               + _vortex[1][0][..., newaxis] * tock(time, period), \
               _vortex[0][1][..., newaxis] * tick(time, period) \
               + _vortex[1][1][..., newaxis] * tock(time, period)
    return _flow

ext: py

max_stars_repo_name: jesserobertson/maxr

max_issues_repo_name: jesserobertson/maxr

max_forks_repo_name: jesserobertson/maxr

reason: from numpy


text:
'''

script for the *mini* Mock Challenge (mini mocha)

>>> python mini_mocha.py arg1 arg2 arg3 arg4 arg5 arg6 arg7 arg8 arg9 arg10 arg11 arg12 arg13 arg14 arg15

* arg1: 'photo' or 'specphoto'
* arg2: simulation (e.g. 'lgal')
* arg3: first galaxy number
* arg4: last galaxy number
* arg5: noise model
    * 'legacy' for `arg1 == photo`
    * 'bgs0_legacy' for `arg1 == specphoto`
* arg6: number of threads
* arg7: number of emcee walkers
* arg8: number of iterations for burnin
* arg9: number of iterations for main chain. If `arg11 == 'adaptive'`, it will
  adaptively determine the number of iterations until convergence
* arg10: maximum number of iterations for `arg11 == adaptive`. This is ignored
  if the number of iteration is specified
* arg11: IF True, overwrites the mcmc file.
* arg12: If True, post processes the MCMC chain.

Here are some examples:

e.g. ifsps run on forward modeled photometry of lgal galaxies 0 to 90 with vanilla
  model. 20 walkers, adaptive mcmc with maximum iteration of 50000. This call will
  overwrite the mcmc file entirely.

>>> python mini_mocha.py photo lgal 0 90 legacy ifsps vanilla 1 20 200 adaptive 50000 overwrite False

e.g. ispeculator run on forward modeled spectrophotometry of lgal galaxies 0 to
  90 with emulator model. 40 walkers, adaptive mcmc with maximum iteration of
  50000. This call will append to existing MCMC chain.

>>> python mini_mocha.py specphoto lgal 0 90 bgs0_legacy ispeculator emulator 1 40 200 adaptive 50000 append False

See `run/cori/lgal_ispeculator_photo.slurm` or
`run/cori/lgal_ispeculator_specphoto.slurm` for examples on how to run this on
NERSC cori.

'''
import sys
import os
import h5py
import numpy as np
import corner as DFM
from functools import partial
from multiprocessing.pool import Pool
from scipy.stats import gaussian_kde as gkde
# --- gqp_mc ---
from gqp_mc import util as UT
from gqp_mc import data as Data
# --- provabgs ---
from provabgs import infer as Infer
from provabgs import flux_calib as FluxCalib
# --- plotting ---
import matplotlib as mpl
import matplotlib.pyplot as plt
try:
    if os.environ['NERSC_HOST'] == 'cori':  pass
except KeyError:
    mpl.rcParams['text.usetex'] = True
mpl.rcParams['font.family'] = 'serif'
mpl.rcParams['axes.linewidth'] = 1.5
mpl.rcParams['axes.xmargin'] = 1
mpl.rcParams['xtick.labelsize'] = 'x-large'
mpl.rcParams['xtick.major.size'] = 5
mpl.rcParams['xtick.major.width'] = 1.5
mpl.rcParams['ytick.labelsize'] = 'x-large'
mpl.rcParams['ytick.major.size'] = 5
mpl.rcParams['ytick.major.width'] = 1.5
mpl.rcParams['legend.frameon'] = False

lbl_dict = {
        'logmstar': r'$\log M_*$',
        'z_metal': r'$\log Z$',
        'dust1': 'dust1',
        'dust2': 'dust2',
        'dust_index': 'dust index',
        'tau': r'$\tau$',
        'f_fiber': r'$f_{\rm fiber}$',
        'beta1_sfh': r'$\beta_1^{\rm SFH}$',
        'beta2_sfh': r'$\beta_2^{\rm SFH}$',
        'beta3_sfh': r'$\beta_3^{\rm SFH}$',
        'beta4_sfh': r'$\beta_4^{\rm SFH}$',
        'gamma1_zh': r'$\gamma_1^{\rm ZH}$',
        'gamma2_zh': r'$\gamma_2^{\rm ZH}$',
        'logsfr.100myr': r'$\log {\rm SFR}_{\rm 100 Myr}$',
        'logsfr.1gyr': r'$\log {\rm SFR}_{\rm 1 Gyr}$',
        'logz.mw': r'$\log Z_{\rm MW}$',
        }

def construct_sample(sim):
    ''' construct the mini Mock Challenge photometry, spectroscopy
    '''
    # construct photometry and spectroscopy
    dir_run = os.path.dirname(os.path.realpath(__file__))
    if sim == 'lgal':
        fscript = os.path.join(dir_run, 'fm_lgal.py')
    elif sim == 'tng':
        fscript = os.path.join(dir_run, 'fm_tng.py')
    os.system('python %s' % fscript)
    return None

def validate_sample(sim):
    ''' generate some plots to validate the mini Mock Challenge photometry
    and spectroscopy
    '''
    # read photometry
    photo, _ = Data.Photometry(sim=sim, noise='legacy', sample='mini_mocha')
    photo_g = 22.5 - 2.5 * np.log10(photo['flux'][:,0])
    photo_r = 22.5 - 2.5 * np.log10(photo['flux'][:,1])
    photo_z = 22.5 - 2.5 * np.log10(photo['flux'][:,2])

    flux_g = photo['flux'][:,0] * 1e-9 * 1e17 * UT.c_light() / 4750.**2 * (3631. * UT.jansky_cgs())
    flux_r = photo['flux'][:,1] * 1e-9 * 1e17 * UT.c_light() / 6350.**2 * (3631. * UT.jansky_cgs())
    flux_z = photo['flux'][:,2] * 1e-9 * 1e17 * UT.c_light() / 9250.**2 * (3631. * UT.jansky_cgs()) # convert to 10^-17 ergs/s/cm^2/Ang
    ivar_g = photo['ivar'][:,0] * (1e-9 * 1e17 * UT.c_light() / 4750.**2 * (3631. * UT.jansky_cgs()))**-2.
    ivar_r = photo['ivar'][:,1] * (1e-9 * 1e17 * UT.c_light() / 6350.**2 * (3631. * UT.jansky_cgs()))**-2.
    ivar_z = photo['ivar'][:,2] * (1e-9 * 1e17 * UT.c_light() / 9250.**2 * (3631. * UT.jansky_cgs()))**-2. # convert to 10^-17 ergs/s/cm^2/Ang

    # read BGS targets from imaging
    bgs_targets = h5py.File(os.path.join(UT.dat_dir(), 'bgs.1400deg2.rlim21.0.hdf5'), 'r')
    n_targets = len(bgs_targets['ra'][...])

    bgs_g = 22.5 - 2.5 * np.log10(bgs_targets['flux_g'][...])[::10]
    bgs_r = 22.5 - 2.5 * np.log10(bgs_targets['flux_r'][...])[::10]
    bgs_z = 22.5 - 2.5 * np.log10(bgs_targets['flux_z'][...])[::10]

    # photometry validation
    fig = plt.figure(figsize=(6,6))
    sub = fig.add_subplot(111)
    DFM.hist2d(bgs_g - bgs_r, bgs_r - bgs_z, color='k', levels=[0.68, 0.95],
            range=[[-1., 3.], [-1., 3.]], bins=40, smooth=0.5,
            plot_datapoints=True, fill_contours=False, plot_density=False, linewidth=0.5,
            ax=sub)
    sub.scatter([-100.], [-100], c='k', s=1, label='BGS targets')
    sub.scatter(photo_g - photo_r, photo_r - photo_z, c='C0', s=1, label='LGal photometry')
    sub.legend(loc='upper left', handletextpad=0, markerscale=10, fontsize=20)
    sub.set_xlabel('$g-r$', fontsize=25)
    sub.set_xlim(-1., 3.)
    sub.set_ylabel('$r-z$', fontsize=25)
    sub.set_ylim(-1., 3.)
    ffig = os.path.join(UT.dat_dir(),
            'mini_mocha', 'mini_mocha.%s.photo.png' % sim)
    fig.savefig(ffig, bbox_inches='tight')
    fig.savefig(UT.fig_tex(ffig, pdf=True), bbox_inches='tight')

    # read spectra
    spec_s, _ = Data.Spectra(sim=sim, noise='none', sample='mini_mocha')
    spec_bgs0, _ = Data.Spectra(sim=sim, noise='bgs0', sample='mini_mocha')
    spec_bgs1, _ = Data.Spectra(sim=sim, noise='bgs1', sample='mini_mocha')
    spec_bgs2, _ = Data.Spectra(sim=sim, noise='bgs2', sample='mini_mocha')

    # read sky brightness
    _fsky = os.path.join(UT.dat_dir(), 'mini_mocha', 'bgs.exposure.surveysim.150s.v0p4.sample.hdf5')
    fsky = h5py.File(_fsky, 'r')
    wave_sky    = fsky['wave'][...] # sky wavelength
    sbright_sky = fsky['sky'][...]

    # spectra validationg
    fig = plt.figure(figsize=(10,5))
    sub = fig.add_subplot(111)

    for i, spec_bgs in enumerate([spec_bgs2, spec_bgs0]):
        wsort = np.argsort(spec_bgs['wave'])
        if i == 0:
            _plt, = sub.plot(spec_bgs['wave'][wsort], spec_bgs['flux'][0,wsort], c='C%i' % i, lw=0.25)
        else:
            sub.plot(spec_bgs['wave'][wsort], spec_bgs['flux'][0,wsort], c='C%i' % i, lw=0.25)

    _plt_photo = sub.errorbar([4750, 6350, 9250], [flux_g[0], flux_r[0], flux_z[0]],
            [ivar_g[0]**-0.5, ivar_r[0]**-0.5, ivar_z[0]**-0.5], fmt='.r')

    _plt_sim, = sub.plot(spec_s['wave'][0,:], spec_s['flux'][0,:], c='k', ls='-', lw=1)
    _plt_sim0, = sub.plot(spec_s['wave'][0,:], spec_s['flux_unscaled'][0,:], c='k', ls=':', lw=0.25)

    leg = sub.legend(
            [_plt_sim0, _plt_photo, _plt_sim, _plt],
            ['%s spectrum' % sim.upper(), '%s photometry' % sim.upper(),
                '%s fiber spectrum' % sim.upper(), '%s BGS spectra' % sim.upper()],
            loc='upper right', fontsize=17)
    for legobj in leg.legendHandles:
        legobj.set_linewidth(2.0)
    sub.set_xlabel('Wavelength [$A$]', fontsize=20)
    sub.set_xlim(3e3, 1e4)
    sub.set_ylabel('flux [$10^{-17} erg/s/cm^2/A$', fontsize=20)
    if sim == 'lgal': sub.set_ylim(-2., 8.)
    elif sim == 'tng': sub.set_ylim(0., None)

    ffig = os.path.join(UT.dat_dir(),
            'mini_mocha', 'mini_mocha.%s.spectra.png' % sim)
    fig.savefig(ffig, bbox_inches='tight')
    fig.savefig(UT.fig_tex(ffig, pdf=True), bbox_inches='tight')

    return None

def postprocess_mcmc(desi_mcmc, fmcmc, dt_sfr=1., n_sample=50000):
    '''  calculate derived properties log M*, log SFR, log Z_MW and prior
    correction weights to impose uniform priors on [log M*, log SFR, log Z_MW]
    '''
    mcmc = desi_mcmc.read_chain(fmcmc)
    flat_chain = desi_mcmc._flatten_chain(mcmc['mcmc_chain'][-1000:,:,:])
    z = mcmc['redshift']

    logmstar = flat_chain[:,0]

    # calculate average SFRs
    avg_sfr_100myr  = desi_mcmc.model.avgSFR(flat_chain, z, dt=0.1)
    avg_sfr_1gyr    = desi_mcmc.model.avgSFR(flat_chain, z, dt=1.)

    avg_logsfr_100myr   = np.log10(avg_sfr_100myr)
    avg_logsfr_1gyr     = np.log10(avg_sfr_1gyr)

    # calculate mass-weighted metallicity
    z_mw    = desi_mcmc.model.Z_MW(flat_chain, z)
    logz_mw = np.log10(z_mw)

    output = {}
    output['logmstar']      = logmstar
    output['logsfr.100myr'] = avg_logsfr_100myr
    output['logsfr.1gyr']   = avg_logsfr_1gyr
    output['logz.mw']       = logz_mw

    # ------------------------------------------------------------
    # get prior correction
    # 1. sample prior
    theta_prior = np.array([desi_mcmc.prior.sample() for i in range(n_sample)])

    # 2. compute the derived properties we want to impose flat priors on
    logm_prior          = theta_prior[:,0]
    logsfr_1gyr_prior   = np.log10(desi_mcmc.model.avgSFR(theta_prior, z, dt=1.))
    logsfr_100myr_prior = np.log10(desi_mcmc.model.avgSFR(theta_prior, z, dt=0.1))
    logzmw_prior    = np.log10(desi_mcmc.model.Z_MW(theta_prior, z))

    # 3. fit a joint distirbution of the derived properties
    kde_fit_1gyr    = gkde(np.array([logm_prior, logsfr_1gyr_prior, logzmw_prior]))
    kde_fit_100myr  = gkde(np.array([logm_prior, logsfr_100myr_prior, logzmw_prior]))

    # 4. calculate the max entropy weights
    p_prop_1gyr = kde_fit_1gyr.pdf(np.array([logmstar, avg_logsfr_1gyr, logz_mw]))
    p_prop_100myr = kde_fit_100myr.pdf(np.array([logmstar, avg_logsfr_100myr, logz_mw]))

    w_prior_1gyr = 1./p_prop_1gyr
    w_prior_100myr = 1./p_prop_100myr

    w_prior_1gyr[p_prop_1gyr < 1e-4] = 0.
    w_prior_100myr[p_prop_100myr < 1e-4] = 0.

    output['w_prior.100myr'] = w_prior_100myr
    output['w_prior.1gyr'] = w_prior_1gyr
    # ------------------------------------------------------------

    for k in ['redshift',
            'wavelength_obs', 'flux_spec_obs', 'flux_ivar_spec_obs',
            'flux_photo_obs', 'flux_ivar_photo_obs',
            'flux_spec_model', 'flux_photo_model']:
        if k in mcmc.keys():
            output[k] = mcmc[k]

    fpost = fmcmc.replace('.mcmc.hdf5', '.postproc.hdf5')
    fh5  = h5py.File(fpost, 'w')
    for k in output.keys():
        if output[k] is None:
            print(k)
            continue
        fh5.create_dataset(k, data=output[k])
    fh5.close()
    return output

def fit_photometry(igal, sim='lgal', noise='legacy', nwalkers=30, burnin=1000,
        niter=1000, maxiter=200000, opt_maxiter=100, overwrite=False,
        postprocess=False):
    ''' Fit Lgal spectra. `noise` specifies whether to fit spectra without
    noise or with BGS-like noise. Produces an MCMC chain and, if not on nersc,
    a corner plot of the posterior.

    :param igal:
        index of Lgal galaxy within the spectral_challenge
    :param noise:
        If 'bgs1'...'bgs8', fit BGS-like spectra. (default: 'none')
    '''
    # read Lgal photometry of the mini_mocha mocks
    photo, meta = Data.Photometry(sim=sim, noise=noise, lib='bc03', sample='mini_mocha')
    assert meta['redshift'][igal] < 0.5, "speculator does not support z > 0.5 yet"

    photo_obs   = photo['flux'][igal,:3]
    ivar_obs    = photo['ivar'][igal,:3]

    print('--- input ---')
    print('  z = %f' % meta['redshift'][igal])
    print('  log M* total = %f' % meta['logM_total'][igal])
    print('  log SFR 100Myr = %f' % np.log10(meta['sfr_100myr'][igal]))
    print('  log SFR 1Gyr = %f' % np.log10(meta['sfr_1gyr'][igal]))
    print('  log Z_MW = %f' % np.log10(meta['Z_MW'][igal]))
    # ------------------------------------------------------------------------------------
    theta_names = ['logmstar', 'beta1_sfh', 'beta2_sfh', 'beta3_sfh', 'beta4_sfh',
            'gamma1_zh', 'gamma2_zh', 'dust1', 'dust2', 'dust_index']
    truths = [None for _ in theta_names]
    truths[0] = meta['logM_total'][igal]

    # set up prior object
    priors = Infer.load_priors([
            Infer.UniformPrior(8, 12, label='sed'),     # uniform priors on logM*
            Infer.FlatDirichletPrior(4, label='sed'),   # flat dirichilet priors
            Infer.UniformPrior(6.9e-5, 7.3e-3, label='sed'),# uniform priors on ZH coeff
            Infer.UniformPrior(6.9e-5, 7.3e-3, label='sed'),# uniform priors on ZH coeff
            Infer.UniformPrior(0., 3., label='sed'),        # uniform priors on dust1
            Infer.UniformPrior(0., 3., label='sed'),        # uniform priors on dust2
            Infer.UniformPrior(-2.2, 0.4, label='sed')      # uniform priors on dust_index
            ])

    desi_mcmc = Infer.desiMCMC(
            prior=priors,
            flux_calib=FluxCalib.no_flux_factor # no flux calibration necessary
            )
    f_mcmc = os.path.join(UT.dat_dir(), 'mini_mocha', 'provabgs',
            '%s.photo.noise_%s.%i.mcmc.hdf5' % (sim, noise, igal))

    if not overwrite and os.path.isfile(f_mcmc):
        print('--- read mcmc ---')
        mcmc = desi_mcmc.read_chain(f_mcmc, debug=True)
    else:
        print('--- run mcmc ---')
        mcmc = desi_mcmc.run(
                photo_obs=photo_obs,
                photo_ivar_obs=ivar_obs,
                zred=meta['redshift'][igal],
                bands='desi',
                sampler='zeus',
                nwalkers=nwalkers,
                burnin=burnin,
                opt_maxiter=opt_maxiter,
                niter=niter,
                writeout=f_mcmc,
                overwrite=overwrite,
                debug=True)

    if postprocess:
        print('--- postprocessing ---')
        post = postprocess_mcmc(desi_mcmc, f_mcmc)

    print('---------------')
    labels = [lbl_dict[_t] for _t in theta_names]

    # corner plot of the posteriors
    flat_chain = desi_mcmc._flatten_chain(mcmc['mcmc_chain'])
    fig = DFM.corner(
        flat_chain,
        range=np.array(mcmc['prior_range']).T,
        quantiles=[0.16, 0.5, 0.84],
        levels=[0.68, 0.95],
        nbin=40,
        smooth=True,
        truths=truths,
        labels=labels,
        label_kwargs={'fontsize': 20})
    fig.savefig(f_mcmc.replace('.hdf5', '.png'), bbox_inches='tight')
    plt.close()

    # corner plot of the posteriors for derive prop
    fig = DFM.corner(
            np.vstack([post['logmstar'], post['logsfr.1gyr'], post['logz.mw']]).T,
            quantiles=[0.16, 0.5, 0.84],
            levels=[0.68, 0.95],
            nbin=40,
            smooth=True,
            hist_kwargs={'density': True})
    _ = DFM.corner(
            np.vstack([post['logmstar'], post['logsfr.1gyr'], post['logz.mw']]).T,
            weights=post['w_prior.1gyr'],
            quantiles=[0.16, 0.5, 0.84],
            levels=[0.68, 0.95],
            nbin=40,
            smooth=True,
            color='C1',
            truths=[meta['logM_total'][igal], np.log10(meta['sfr_1gyr'][igal]),
                np.log10(meta['Z_MW'][igal])],
            labels=[lbl_dict['logmstar'], lbl_dict['logsfr.1gyr'], lbl_dict['logz.mw']],
            label_kwargs={'fontsize': 20},
            fig=fig,
            hist_kwargs={'density': True})
    fig.savefig(f_mcmc.replace('.hdf5', '.prop.png'), bbox_inches='tight')
    plt.close()

    # best-fit observable
    fig = plt.figure(figsize=(5,3))
    sub = fig.add_subplot(111)
    sub.errorbar(np.arange(len(photo_obs)), photo_obs,
            yerr=ivar_obs**-0.5, fmt='.k', label='data')
    sub.scatter(np.arange(len(photo_obs)), mcmc['flux_photo_model'], c='C1',
            label='best-fit')
    sub.legend(loc='upper left', markerscale=3, handletextpad=0.2, fontsize=15)
    sub.set_xticks([0, 1, 2])
    sub.set_xticklabels(['$g$', '$r$', '$z$'])
    sub.set_xlim(-0.5, len(photo_obs)-0.5)
    fig.savefig(f_mcmc.replace('.hdf5', '.bestfit.png'), bbox_inches='tight')
    plt.close()
    return None

def fit_spectrum(igal, sim='lgal', noise='bgs', nwalkers=30, burnin=1000,
        niter=1000, maxiter=200000, opt_maxiter=100, overwrite=False,
        postprocess=False):
    ''' Fit Lgal spectra. `noise` specifies whether to fit spectra without
    noise or with BGS-like noise. Produces an MCMC chain and, if not on nersc,
    a corner plot of the posterior.

    :param igal:
        index of Lgal galaxy within the spectral_challenge
    :param noise:
        If 'bgs1'...'bgs8', fit BGS-like spectra. (default: 'none')
    '''
    # read Lgal spectra of the spectral_challenge mocks
    specs, meta = Data.Spectra(sim=sim, noise=noise, lib='bc03', sample='mini_mocha')

    assert meta['redshift'][igal] < 0.5, "speculator does not support z > 0.5 yet"

    wave_obs        = specs['wave']
    print(wave_obs.shape)
    flux_obs        = specs['flux'][igal]
    if noise != 'none':
        ivar_obs    = specs['ivar'][igal]
        resolution  = [specs['res_b'][igal], specs['res_r'][igal], specs['res_z'][igal]]
    else:
        ivar_obs    = np.ones(len(specs['wave']))
        resolution  = None

    print('--- input ---')
    print('  z = %f' % meta['redshift'][igal])
    print('  log M* fiber = %f' % meta['logM_fiber'][igal])
    print('  log SFR 100Myr = %f' % np.log10(meta['sfr_100myr'][igal]))
    print('  log SFR 1Gyr = %f' % np.log10(meta['sfr_1gyr'][igal]))
    print('  log Z_MW = %f' % np.log10(meta['Z_MW'][igal]))
    # ------------------------------------------------------------------------------------
    theta_names = ['logmstar', 'beta1_sfh', 'beta2_sfh', 'beta3_sfh', 'beta4_sfh',
            'gamma1_zh', 'gamma2_zh', 'dust1', 'dust2', 'dust_index']
    truths = [None for _ in theta_names]
    truths[0] = meta['logM_fiber'][igal]

    # set up prior object
    priors = Infer.load_priors([
            Infer.UniformPrior(8, 12, label='sed'),     # uniform priors on logM*
            Infer.FlatDirichletPrior(4, label='sed'),   # flat dirichilet priors
            Infer.UniformPrior(6.9e-5, 7.3e-3, label='sed'),# uniform priors on ZH coeff
            Infer.UniformPrior(6.9e-5, 7.3e-3, label='sed'),# uniform priors on ZH coeff
            Infer.UniformPrior(0., 3., label='sed'),        # uniform priors on dust1
            Infer.UniformPrior(0., 3., label='sed'),        # uniform priors on dust2
            Infer.UniformPrior(-2.2, 0.4, label='sed')      # uniform priors on dust_index
            ])

    desi_mcmc = Infer.desiMCMC(
            prior=priors,
            flux_calib=FluxCalib.no_flux_factor # no flux calibration necessary
            )
    f_mcmc = os.path.join(UT.dat_dir(), 'mini_mocha', 'provabgs',
            '%s.spec.noise_%s.%i.mcmc.hdf5' % (sim, noise, igal))

    if not overwrite and os.path.isfile(f_mcmc):
        print('--- read mcmc ---')
        mcmc = desi_mcmc.read_chain(f_mcmc, debug=True)
    else:
        print('--- run mcmc ---')
        mcmc = desi_mcmc.run(
                wave_obs=wave_obs,
                flux_obs=flux_obs,
                flux_ivar_obs=ivar_obs,
                resolution=resolution,
                zred=meta['redshift'][igal],
                vdisp=0., # km/s velocity dispersion
                mask='emline',
                sampler='zeus',
                nwalkers=nwalkers,
                burnin=burnin,
                opt_maxiter=opt_maxiter,
                niter=niter,
                writeout=f_mcmc,
                overwrite=overwrite,
                debug=True)

    if postprocess:
        print('--- postprocessing ---')
        post = postprocess_mcmc(desi_mcmc, f_mcmc)

    print('---------------')
    labels = [lbl_dict[_t] for _t in theta_names]

    # corner plot of the posteriors
    flat_chain = desi_mcmc._flatten_chain(mcmc['mcmc_chain'])
    fig = DFM.corner(
            flat_chain[:-3],
        range=np.array(mcmc['prior_range']).T,
        quantiles=[0.16, 0.5, 0.84],
        levels=[0.68, 0.95],
        nbin=40,
        smooth=True,
        truths=truths,
        labels=labels,
        label_kwargs={'fontsize': 20})
    fig.savefig(f_mcmc.replace('.hdf5', '.png'), bbox_inches='tight')
    plt.close()

    # corner plot of the posteriors for derive prop
    fig = DFM.corner(
            np.vstack([post['logmstar'], post['logsfr.1gyr'], post['logz.mw']]).T,
            quantiles=[0.16, 0.5, 0.84],
            levels=[0.68, 0.95],
            nbin=40,
            smooth=True,
            hist_kwargs={'density': True})
    _ = DFM.corner(
            np.vstack([post['logmstar'], post['logsfr.1gyr'], post['logz.mw']]).T,
            weights=post['w_prior.1gyr'],
            quantiles=[0.16, 0.5, 0.84],
            levels=[0.68, 0.95],
            nbin=40,
            smooth=True,
            color='C1',
            truths=[meta['logM_fiber'][igal], np.log10(meta['sfr_1gyr'][igal]),
                np.log10(meta['Z_MW'][igal])],
            labels=[lbl_dict['logmstar'], lbl_dict['logsfr.1gyr'], lbl_dict['logz.mw']],
            label_kwargs={'fontsize': 20},
            fig=fig,
            hist_kwargs={'density': True})
    fig.savefig(f_mcmc.replace('.hdf5', '.prop.png'), bbox_inches='tight')
    plt.close()

    # best-fit observable
    fig = plt.figure(figsize=(15,5))
    sub = fig.add_subplot(111)
    sub.plot(mcmc['wavelength_obs'], mcmc['flux_spec_obs'],
        c='C0', lw=1, label='obs.')
    sub.plot(mcmc['wavelength_obs'], mcmc['flux_spec_model'],
        c='k', ls='--', lw=1, label='best-fit')
    sub.legend(loc='upper right', fontsize=15)
    sub.set_xlabel('wavelength [$A$]', fontsize=20)
    sub.set_xlim(3600., 9800.)
    sub.set_ylim(-1., 5.)
    fig.savefig(f_mcmc.replace('.hdf5', '.bestfit.png'), bbox_inches='tight')
    plt.close()
    return None

def fit_spectrophotometry(igal, sim='lgal', noise='bgs0_legacy', nwalkers=30,
        burnin=1000, niter=1000, maxiter=200000, opt_maxiter=100,
        overwrite=False, postprocess=False):
    ''' Fit Lgal spectra. `noise` specifies whether to fit spectra without
    noise or with BGS-like noise. Produces an MCMC chain and, if not on nersc,
    a corner plot of the posterior.

    :param igal:
        index of Lgal galaxy within the spectral_challenge
    :param noise:
        If 'bgs1'...'bgs8', fit BGS-like spectra. (default: 'none')
    '''
    noise_spec = noise.split('_')[0]
    noise_photo = noise.split('_')[1]
    # read noiseless Lgal spectra of the spectral_challenge mocks
    specs, meta = Data.Spectra(sim=sim, noise=noise_spec, lib='bc03', sample='mini_mocha')
    # read Lgal photometry of the mini_mocha mocks
    photo, _ = Data.Photometry(sim=sim, noise=noise_photo, lib='bc03', sample='mini_mocha')

    assert meta['redshift'][igal] < 0.5, "speculator does not support z > 0.5 yet"

    wave_obs        = specs['wave']
    flux_obs        = specs['flux'][igal]
    ivar_obs        = specs['ivar'][igal]
    resolution      = [specs['res_b'][igal], specs['res_r'][igal], specs['res_z'][igal]]
    photo_obs       = photo['flux'][igal,:3]
    photo_ivar_obs  = photo['ivar'][igal,:3]

    # get fiber flux factor prior range based on "measured" fiber flux
    f_fiber_true = (photo['fiberflux_r_true'][igal]/photo['flux_r_true'][igal])
    prior_width = np.max([0.05, 5.*photo['fiberflux_r_ivar'][igal]**-0.5/photo['flux'][igal,1]])
    f_fiber_min = (photo['fiberflux_r_meas'][igal])/photo['flux'][igal,1] - prior_width
    f_fiber_max = (photo['fiberflux_r_meas'][igal])/photo['flux'][igal,1] + prior_width

    print('--- input ---')
    print('  z = %f' % meta['redshift'][igal])
    print('  log M* total = %f' % meta['logM_total'][igal])
    print('  log M* fiber = %f' % meta['logM_fiber'][igal])
    print('  log SFR 100Myr = %f' % np.log10(meta['sfr_100myr'][igal]))
    print('  log SFR 1Gyr = %f' % np.log10(meta['sfr_1gyr'][igal]))
    print('  log Z_MW = %f' % np.log10(meta['Z_MW'][igal]))
    print('  f_fiber = %f' % f_fiber_true)
    print('  f_fiber range = [%f, %f]' % (f_fiber_min, f_fiber_max))
    # ------------------------------------------------------------------------------------
    theta_names = ['logmstar', 'beta1_sfh', 'beta2_sfh', 'beta3_sfh', 'beta4_sfh',
            'gamma1_zh', 'gamma2_zh', 'dust1', 'dust2', 'dust_index', 'f_fiber']
    truths = [None for _ in theta_names]
    truths[0] = meta['logM_total'][igal]
    truths[-1] = f_fiber_true

    # set up prior object
    priors = Infer.load_priors([
            Infer.UniformPrior(8, 12, label='sed'),     # uniform priors on logM*
            Infer.FlatDirichletPrior(4, label='sed'),   # flat dirichilet priors
            Infer.UniformPrior(6.9e-5, 7.3e-3, label='sed'),# uniform priors on ZH coeff
            Infer.UniformPrior(6.9e-5, 7.3e-3, label='sed'),# uniform priors on ZH coeff
            Infer.UniformPrior(0., 3., label='sed'),        # uniform priors on dust1
            Infer.UniformPrior(0., 3., label='sed'),        # uniform priors on dust2
            Infer.UniformPrior(-2.2, 0.4, label='sed'),     # uniform priors on dust_index
            Infer.UniformPrior(f_fiber_min, f_fiber_max, label='flux_calib')
            ])

    desi_mcmc = Infer.desiMCMC(
            prior=priors,
            flux_calib=FluxCalib.constant_flux_factor # no flux calibration necessary
            )
    f_mcmc = os.path.join(UT.dat_dir(), 'mini_mocha', 'provabgs',
            '%s.specphoto.noise_%s.%i.mcmc.hdf5' % (sim, noise, igal))

    if not overwrite and os.path.isfile(f_mcmc):
        print('--- read mcmc ---')
        mcmc = desi_mcmc.read_chain(f_mcmc, debug=True)
    else:
        print('--- run mcmc ---')
        mcmc = desi_mcmc.run(
                wave_obs=wave_obs,
                flux_obs=flux_obs,
                flux_ivar_obs=ivar_obs,
                resolution=resolution,
                photo_obs=photo_obs,
                photo_ivar_obs=photo_ivar_obs,
                zred=meta['redshift'][igal],
                vdisp=0.,
                mask='emline',
                bands='desi',
                sampler='zeus',
                nwalkers=nwalkers,
                burnin=burnin,
                opt_maxiter=opt_maxiter,
                niter=niter,
                writeout=f_mcmc,
                overwrite=overwrite,
                debug=True)

    if postprocess:
        print('--- postprocessing ---')
        post = postprocess_mcmc(desi_mcmc, f_mcmc)

    print('---------------')

    # corner plot of the posteriors
    flat_chain = desi_mcmc._flatten_chain(mcmc['mcmc_chain'])
    fig = DFM.corner(
        flat_chain,
        range=np.array(mcmc['prior_range']).T,
        quantiles=[0.16, 0.5, 0.84],
        levels=[0.68, 0.95],
        nbin=40,
        smooth=True,
        truths=truths,
        labels=[lbl_dict[_t] for _t in theta_names],
        label_kwargs={'fontsize': 20})
    fig.savefig(f_mcmc.replace('.hdf5', '.png'), bbox_inches='tight')
    plt.close()

    # corner plot of the posteriors for derive prop
    fig = DFM.corner(
            np.vstack([post['logmstar'], post['logsfr.1gyr'], post['logz.mw']]).T,
            quantiles=[0.16, 0.5, 0.84],
            levels=[0.68, 0.95],
            nbin=40,
            smooth=True,
            hist_kwargs={'density': True})
    _ = DFM.corner(
            np.vstack([post['logmstar'], post['logsfr.1gyr'], post['logz.mw']]).T,
            weights=post['w_prior.1gyr'],
            quantiles=[0.16, 0.5, 0.84],
            levels=[0.68, 0.95],
            nbin=40,
            smooth=True,
            color='C1',
            truths=[meta['logM_total'][igal], np.log10(meta['sfr_1gyr'][igal]),
                np.log10(meta['Z_MW'][igal])],
            labels=[lbl_dict['logmstar'], lbl_dict['logsfr.1gyr'], lbl_dict['logz.mw']],
            label_kwargs={'fontsize': 20},
            fig=fig,
            hist_kwargs={'density': True})
    fig.savefig(f_mcmc.replace('.hdf5', '.prop.png'), bbox_inches='tight')
    plt.close()

    # best-fit observable
    fig = plt.figure(figsize=(18,5))
    gs = mpl.gridspec.GridSpec(1, 2, figure=fig, width_ratios=[1,3], wspace=0.2)
    sub = plt.subplot(gs[0])
    sub.errorbar(np.arange(len(photo_obs)), photo_obs,
            yerr=photo_ivar_obs**-0.5, fmt='.k', label='data')
    sub.scatter(np.arange(len(photo_obs)), mcmc['flux_photo_model'], c='C1',
            label='model')
    sub.legend(loc='upper left', markerscale=3, handletextpad=0.2, fontsize=15)
    sub.set_xticks([0, 1, 2])
    sub.set_xticklabels(['$g$', '$r$', '$z$'], fontsize=25)
    sub.set_xlim(-0.5, len(photo_obs)-0.5)

    sub = plt.subplot(gs[1])
    sub.plot(mcmc['wavelength_obs'], mcmc['flux_spec_obs'],
        c='C0', lw=1, label='obs.')
    sub.plot(mcmc['wavelength_obs'], mcmc['flux_spec_model'],
        c='k', ls='--', lw=1, label='best-fit')
    sub.legend(loc='upper right', fontsize=15)
    sub.set_xlabel('wavelength [$A$]', fontsize=20)
    sub.set_xlim(3600., 9800.)
    sub.set_ylim(-1., 5.)
    fig.savefig(f_mcmc.replace('.hdf5', '.bestfit.png'), bbox_inches='tight')
    plt.close()
    return None

def MP_sed_fit(spec_or_photo, igals, sim='lgal', noise='none', nthreads=1,
        nwalkers=100, burnin=100, niter=1000, maxiter=200000, overwrite=False,
        postprocess=False):
    ''' multiprocessing wrapepr for fit_spectra and fit_photometry. This does *not* parallelize
    the MCMC sampling of individual fits but rather runs multiple fits simultaneously.

    :param spec_or_photo:
        fit spectra or photometry

    :param igals:
        array/list of spectral_challenge galaxy indices

    :param noise:
        If 'none', fit noiseless spectra.
        If 'bgs1'...'bgs8', fit BGS-like spectra. (default: 'none')

    :param dust:
        If True, fit the spectra w/ dust using a model with dust
        If False, fit the spectra w/o dust using a model without dust.
        (default: False)

    :param nthreads:
        Number of threads. If nthreads == 1, just runs fit_spectra
    '''
    args = igals # galaxy indices

    kwargs = {
            'sim': sim,
            'noise': noise,
            'nwalkers': nwalkers,
            'burnin': burnin,
            'niter': niter,
            'maxiter': maxiter,
            'opt_maxiter': 10000,
            'overwrite': overwrite,
            'postprocess': postprocess
            }
    if spec_or_photo == 'spec':
        fit_func = fit_spectrum
    elif spec_or_photo == 'photo':
        fit_func = fit_photometry
    elif spec_or_photo == 'specphoto':
        fit_func = fit_spectrophotometry

    if nthreads > 1:
        pool = Pool(processes=nthreads)
        pool.map(partial(fit_func, **kwargs), args)
        pool.close()
        pool.terminate()
        pool.join()
    else:
        # single thread, loop over
        for igal in args: fit_func(igal, **kwargs)
    return None

if __name__=="__main__":
    # >>> python mini_mocha.py
    spec_or_photo   = sys.argv[1]
    sim             = sys.argv[2]

    if spec_or_photo == 'construct':
        construct_sample(sim)
        validate_sample(sim)
        sys.exit()

    igal0           = int(sys.argv[3])
    igal1           = int(sys.argv[4])
    noise           = sys.argv[5]
    nthreads        = int(sys.argv[6])
    nwalkers        = int(sys.argv[7])
    burnin          = int(sys.argv[8])
    niter           = sys.argv[9]
    maxiter         = int(sys.argv[10])
    overwrite       = sys.argv[11] == 'True'
    postprocess     = sys.argv[12] == 'True'

    try:
        niter = int(niter)
    except ValueError:
        pass

    print('----------------------------------------')
    print('fitting %s of mini_mocha galaxies %i to %i' %
            (spec_or_photo, igal0, igal1))
    print('using %i threads' % nthreads)
    igals = range(igal0, igal1+1)

    MP_sed_fit(spec_or_photo,
            igals,
            sim=sim,
            noise=noise,
            nthreads=nthreads,
            nwalkers=nwalkers,
            burnin=burnin,
            niter=niter,
            maxiter=maxiter,
            overwrite=overwrite,
            postprocess=postprocess)

reason: import numpy,from scipy


text:
import os

import numpy as np

import gradio as gr

def reverse_audio(audio):
    sr, data = audio
    return (sr, np.flipud(data))

demo = gr.Interface(fn=reverse_audio,
                    inputs="microphone",
                    outputs="audio",
                    examples=os.path.join(os.path.dirname(__file__), "audio"))

if __name__ == "__main__":
    demo.launch()

max_stars_repo_name: dumpmemory/gradio

reason: import numpy


text:
import itertools
import numpy as np

class agentens:

    def __init__(self, setup : dict, **kwargs):
        """
        A class for an ensemble of agents which
        can be evolved in time.

        Args:
        setup
            A dictionary of setup parameters.

        """
        self.setup = setup
        self.tau = None

        # Calculate the total number of agents and store it
        self.setup["Nagents"] = (
            self.setup["Nproviders"] + self.setup["Ntakers"]
        )

        # Setup the bid and ask decision properties
        self.bids = np.zeros(
            (self.setup["Nlattice"], self.setup["Nproviders"]),
            dtype=int,
        )
        self.asks = np.zeros(
            (self.setup["Nlattice"], self.setup["Nproviders"]),
            dtype=int,
        )

        # Keep a memory of all outstanding limit orders,
        # their volumes and which agent they are associated to
        self.membidLOs = np.zeros(
            (self.setup["Nlattice"], self.setup["Nproviders"]),
            dtype=int,
        )
        self.memaskLOs = np.zeros(
            (self.setup["Nlattice"], self.setup["Nproviders"]),
            dtype=int,
        )

        # Draw the initial latent market order volumes and signs
        # from the specified Pareto distribution
        self.latmovs = (
            np.random.pareto(
                self.setup["MOvolpower"],
                size=self.setup["Ntakers"],
            ) + 1.0
        ).astype(int)
        self.latmovs[
            self.latmovs > self.setup["MOvolcutoff"]
        ] = self.setup["MOvolcutoff"]
        self.mosigns = (
            1.0 - (
                2.0
                * np.random.binomial(
                    1,
                    0.5,
                    size=self.setup["Ntakers"]
                )
            )
        )

        # Draw the trader agressiveness factors for market
        # and limit orders from a beta distribution
        self.moaggrf = np.random.beta(
            self.setup["MOaggrA"],
            self.setup["MOaggrB"],
            size=self.setup["Ntakers"],
        )
        self.loaggrf = np.random.beta(
            self.setup["LOaggrA"],
            self.setup["LOaggrB"],
            size=self.setup["Nproviders"],
        )

    def iterate(self, market_state_info : dict):
        """
        Iterate the ensemble a step forward in time by
        asking each agent to make buy-sell-cancel-hold decisions.

        Args:
        market_state_info
            A dictionary of current market state info.

        """
        # Sum over past limit orders by agent
        summembidLOs = np.sum(self.membidLOs, axis=0)
        summemaskLOs = np.sum(self.memaskLOs, axis=0)

        # Consistent event rate computations with
        # agent speculation taken into account
        self.tau = np.random.exponential(1.0 / self.setup["meanHOrate"])
        HOr, LOrb, LOra, MOrb, MOra, COrb, COra = (
            (1.0 / self.tau),
            self.setup["meanLOratebid"] * (
                1.0 - (
                    self.setup["LOsignalpha"]
                    * market_state_info["bidptdrop"]
                    * np.ones(self.setup["Nproviders"])
                )
            ),
            self.setup["meanLOrateask"] * (
                1.0 + (
                    self.setup["LOsignalpha"]
                    * market_state_info["askptrise"]
                    * np.ones(self.setup["Nproviders"])
                )
            ),
            self.setup["meanMOratebid"] * (1.0 + self.mosigns),
            self.setup["meanMOrateask"] * (1.0 - self.mosigns),
            summembidLOs * self.setup["meanCOratebid"],
            summemaskLOs * self.setup["meanCOrateask"],
        )
        provtotr = HOr + LOrb + LOra + COrb + COra
        taketotr = HOr + MOrb + MOra

        # Draw events from the uniform distribution
        # and apply the rejection inequalities to
        # assign the agents' decisions
        evs = np.random.uniform(size=self.setup["Nagents"])
        provevs = evs[:self.setup["Nproviders"]]
        takeevs = evs[self.setup["Nproviders"]:]
        LOsb = (provevs < LOrb / provtotr)
        LOsa = (
            (LOrb / provtotr <= provevs)
            * (provevs < (LOra + LOrb) / provtotr)
        )
        MOsb = (takeevs < MOrb / taketotr)
        MOsa = (
            (MOrb / taketotr <= takeevs)
            & (takeevs < (MOrb + MOra) / taketotr)
        )
        COsb = (
            ((LOra + LOrb) / provtotr <= provevs)
            & (provevs < (LOra + LOrb + COrb) / provtotr)
        )
        COsa = (
            ((LOra + LOrb + COrb) / provtotr <= provevs)
            & (
                provevs < (
                    LOra + LOrb + COrb + COra
                ) / provtotr
            )
        )

        # Decide on the prices for both the limit orders
        # and the order cancellations
        provboa = (LOsb | COsb)
        prs = np.random.uniform(
            size=(self.setup["Nlattice"], self.setup["Nproviders"])
        )
        dec = np.exp(
            - np.abs(
                market_state_info["midprice"]
                - market_state_info["prices"]
            ) * self.setup["meanLOdecay"]
        )
        midpt = float(
            market_state_info["bidpt"] + market_state_info["askpt"]
        ) / 2.0
        midptlow = int(np.floor(midpt))
        midpthigh = int(np.ceil(midpt))
        LObpts = np.random.choice(
            np.arange(0, midptlow + 1, 1, dtype=int),
            size=self.setup["Nproviders"],
            p=(
                dec[:midptlow + 1] / np.sum(dec[:midptlow + 1])
            ),
        )
        LOapts = np.random.choice(
            np.arange(midpthigh, self.setup["Nlattice"], 1, dtype=int),
            size=self.setup["Nproviders"],
            p=(
                dec[midpthigh:] / np.sum(dec[midpthigh:])
            ),
        )
        CObpts = np.argmax(
            (
                prs
                * ((self.membidLOs + self.memaskLOs) > 0)
                * provboa
            ),
            axis=0,
        )
        COapts = np.argmax(
            (
                prs
                * ((self.membidLOs + self.memaskLOs) > 0)
                * (provboa==False)
            ),
            axis=0,
        )

        # Pass the limit-order and market-order decisions
        # on to the output properties while leaving some
        # market orders unfulfilled if there are too many
        # of them
        self.bids[:], self.asks[:] = 0, 0
        naglosa = np.sum(self.memaskLOs[market_state_info["askpt"]])
        naglosb = np.sum(self.membidLOs[market_state_info["bidpt"]])
        LObvols = (self.loaggrf[LOsb] * naglosb).astype(int)
        LOavols = (self.loaggrf[LOsa] * naglosa).astype(int)
        LObvols[LObvols==0] = 1
        LOavols[LOavols==0] = 1
        self.bids[
            (
                LObpts[LOsb],
                np.arange(0, self.setup["Nproviders"], 1, dtype=int)[
                    LOsb
                ],
            )
        ] += LObvols
        self.asks[
            (
                LOapts[LOsa],
                np.arange(0, self.setup["Nproviders"], 1, dtype=int)[
                    LOsa
                ],
            )
        ] += LOavols
        nMOsa, nMOsb = int(np.sum(MOsa)), int(np.sum(MOsb))
        if nMOsb > 0:
            aglos = np.random.permutation(
                np.repeat(
                    np.arange(0, self.setup["Nproviders"], 1, dtype=int)[
                        self.memaskLOs[market_state_info["askpt"]] > 0
                    ],
                    self.memaskLOs[market_state_info["askpt"]][
                        self.memaskLOs[market_state_info["askpt"]] > 0
                    ]
                )
            )
            agmos = np.random.choice(
                np.arange(0, self.setup["Ntakers"], 1, dtype=int)[MOsb],
                size=nMOsb,
                replace=False,
            )
            exsize = (self.moaggrf[agmos] * naglosa).astype(int)
            exsize[exsize==0] = 1
            exsize[exsize > self.latmovs[agmos]] = self.latmovs[agmos][
                exsize > self.latmovs[agmos]
            ]
            cexsize = np.cumsum(exsize)
            fullyexmask = (cexsize < naglosa)
            partexind = len(cexsize[fullyexmask])
            upto = 0
            if np.any(fullyexmask):
                upto = int(np.max(cexsize[fullyexmask]))
                upto += (len(cexsize) > partexind) * (naglosa - upto)
            self.asks[market_state_info["askpt"]][aglos[:upto]] -= 1
            self.latmovs[agmos[fullyexmask]] -= exsize[fullyexmask]
            if len(cexsize) > partexind:
                self.latmovs[agmos[partexind]] -= (naglosa - upto)
            nchanges = int(np.sum(self.latmovs==0))
            self.mosigns[self.latmovs==0] = (
                1.0 - (
                    2.0 * np.random.binomial(1, 0.5, size=nchanges)
                )
            )
            self.latmovs[self.latmovs==0] = (
                np.random.pareto(
                    self.setup["MOvolpower"],
                    size=nchanges,
                ) + 1.0
            ).astype(int)
            self.latmovs[
                self.latmovs > self.setup["MOvolcutoff"]
            ] = self.setup["MOvolcutoff"]
        if nMOsa > 0:
            aglos = np.random.permutation(
                np.repeat(
                    np.arange(0, self.setup["Nproviders"], 1, dtype=int)[
                        self.membidLOs[market_state_info["bidpt"]] > 0
                    ],
                    self.membidLOs[market_state_info["bidpt"]][
                        self.membidLOs[market_state_info["bidpt"]] > 0
                    ]
                )
            )
            agmos = np.random.choice(
                np.arange(0, self.setup["Ntakers"], 1, dtype=int)[MOsa],
                size=nMOsa,
                replace=False,
            )
            exsize = (self.moaggrf[agmos] * naglosb).astype(int)
            exsize[exsize==0] = 1
            exsize[exsize > self.latmovs[agmos]] = self.latmovs[agmos][
                exsize > self.latmovs[agmos]
            ]
            cexsize = np.cumsum(exsize)
            fullyexmask = (cexsize < naglosb)
            partexind = len(cexsize[fullyexmask])
            upto = 0
            if np.any(fullyexmask):
                upto = int(np.max(cexsize[fullyexmask]))
                upto += (len(cexsize) > partexind) * (naglosb - upto)
            self.bids[market_state_info["bidpt"]][aglos[:upto]] -= 1
            self.latmovs[agmos[fullyexmask]] -= exsize[fullyexmask]
            if len(cexsize) > partexind:
                self.latmovs[agmos[partexind]] -= (naglosb - upto)
            nchanges = int(np.sum(self.latmovs==0))
            self.mosigns[self.latmovs==0] = (
                1.0 - (
                    2.0 * np.random.binomial(1, 0.5, size=nchanges)
                )
            )
            self.latmovs[self.latmovs==0] = (
                np.random.pareto(
                    self.setup["MOvolpower"],
                    size=nchanges,
                ) + 1.0
            ).astype(int)
            self.latmovs[
                self.latmovs > self.setup["MOvolcutoff"]
            ] = self.setup["MOvolcutoff"]

        # Pass the cancel-order decisions on to the output
        # properties if they haven't already been fulfilled
        nbidlos = np.sum(self.membidLOs + self.bids, axis=1)
        nasklos = np.sum(self.memaskLOs + self.asks, axis=1)
        cbids = np.zeros(
            (self.setup["Nlattice"], self.setup["Nproviders"]),
        )
        casks = np.zeros(
            (self.setup["Nlattice"], self.setup["Nproviders"]),
        )
        cbids[
            (
                CObpts[COsb],
                np.arange(0, self.setup["Nproviders"], 1, dtype=int)[
                    COsb
                ],
            )
        ] += np.random.uniform(size=len(CObpts[COsb]))
        casks[
            (
                COapts[COsa],
                np.arange(0, self.setup["Nproviders"], 1, dtype=int)[
                    COsa
                ],
            )
        ] += np.random.uniform(size=len(COapts[COsa]))
        cbidssinds = np.argsort(cbids, axis=1)
        caskssinds = np.argsort(casks, axis=1)
        cbput = 0 + 1 * (
            np.tensordot(
                self.setup["Nlattice"] - np.minimum(
                    np.sum(cbids > 0, axis=1),
                    nbidlos,
                ),
                np.ones(self.setup["Nproviders"]),
                axes=0,
            ) <= np.tensordot(
                np.ones(self.setup["Nlattice"]),
                np.arange(0, self.setup["Nproviders"], 1),
                axes=0,
            )
        )
        caput = 0 + 1 * (
            np.tensordot(
                self.setup["Nlattice"] - np.minimum(
                    np.sum(casks > 0, axis=1),
                    nasklos,
                ),
                np.ones(self.setup["Nproviders"]),
                axes=0,
            ) <= np.tensordot(
                np.ones(self.setup["Nlattice"]),
                np.arange(0, self.setup["Nproviders"], 1),
                axes=0,
            )
        )
        np.put_along_axis(cbids, cbidssinds, cbput, axis=1)
        np.put_along_axis(casks, caskssinds, caput, axis=1)
        self.bids -= cbids.astype(int)
        self.asks -= casks.astype(int)

        # Update the agent-specific limit order memory
        self.membidLOs += self.bids
        self.memaskLOs += self.asks

max_stars_repo_name: umbralcalc/lobsim

max_issues_repo_name: umbralcalc/lobsim

max_forks_repo_name: umbralcalc/lobsim


text:
"""
Functions for defining the effect of Milky Way dust on SAMI observations.

For a variety of published maps (currently 2) this module will download the
maps and use them to measure the E(B-V) value at the position of a SAMI
observation. Then for different parameterisations of the dust attenuation
law it will calculate the transmission curve for the observation. The
results are saved in the FITS file, but no correction is made to the data
itself.
"""
from __future__ import absolute_import, division, print_function, unicode_literals

import os
import numpy as np
from math import pi, sqrt, sin, cos
from matplotlib import pyplot as plt

try:
    import urllib.request as urllib2
except ImportError:
    import urllib2

import astropy.io.fits as pf

try :
    import healpy as hp
    HEALPY_AVAILABLE = True
except ImportError:
    HEALPY_AVAILABLE = False
    # print('Requires healpy -- developed with py27-healpy @1.8.1_0')
    # print('If you use Macports to manage your python installation,')
    # print('type "sudo port install py27-healpy".')

from astropy import coordinates as co
from astropy import units as u

MAPS = {}

MAPS_FILES = {
    'planck': {
        'filename': 'dust/HFI_CompMap_ThermalDustModel_2048_R1.20.fits',
        'field': 2,
        'url': 'http://pla.esac.esa.int/pla/aio/product-action?MAP.MAP_ID=HFI_CompMap_ThermalDustModel_2048_R1.20.fits',
        'header_key': 'EBVPLNCK',
        'comment_name': 'Planck v1.20'},
    'sfd98': {
        'filename': 'dust/lambda_sfd_ebv.fits',
        'field': 0,
        'url': 'http://lambda.gsfc.nasa.gov/data/foregrounds/SFD/lambda_sfd_ebv.fits',
        'header_key': 'EBVSFD98',
        'comment_name': 'SFD98'},
    # 'schlafly': {
    #     'filename': 'dust/ps1-ebv-4.5kpc.fits',
    #     'field': 0,
    #     'url': 'http://lambda.gsfc.nasa.gov/data/foregrounds/EBV/ps1-ebv-4.5kpc.fits',
    #     'header_key': 'EBVSCHLA',
    #     'comment_name': 'Schlafly et al 2014'},
}

def load_map(name, force_reload=False):
    """Load the dust maps from various sources."""
    if name not in MAPS or force_reload:
        map_info = MAPS_FILES[name]
        try:
            MAPS[name] = hp.read_map(
                map_info['filename'], field=map_info['field'])
        except IOError:
            return False
    return True

def download_map(name, overwrite=False):
    """Download a single dust map."""
    map_info = MAPS_FILES[name]
    if os.path.exists(map_info['filename']) and not overwrite:
        print('{} map already downloaded; returning')
        return
    dirname = os.path.dirname(map_info['filename'])
    if not os.path.exists(dirname):
        os.makedirs(dirname)
    response = urllib2.urlopen(map_info['url'])
    with open(map_info['filename'], 'w') as f_out:
        f_out.write(response.read())

def download_all_maps(overwrite=False):
    """Download all the dust maps."""
    for name in MAPS_FILES:
        download_map(name, overwrite=overwrite)

# planck_dust_map_filename = 'HFI_CompMap_ThermalDustModel_2048_R1.20.fits'
# if not os.path.exists( planck_dust_map_filename ):
#     print('WARNING: Cannot find Planck dust map with file name:')
#     print(' '*8, planck_dust_map_filename)
#     print()
#     print('This can be downloaded via the Planck Explanatory Supplement wiki:')
#     print(' '*8, 'http://wiki.cosmos.esa.int/planckpla')
#     print('(Look under Mission Products > CMB and astrophysical component maps,')
#     print('and be sure to download the higher resolution Nside=2048 version.')
#     print('Note that this is a ~1.6 Gb download!)')
#     print()
# else :
#     print('reading Planck dust map from file:')
#     print(' '*8, planck_dust_map_filename)
#     planckDustMap = hp.read_map( planck_dust_map_filename, field=2 )
#     hp.mollview( planckDustMap, min=0., max=1., fig=0,
#                  title='Planck dust map: %s'
#                  % planck_dust_map_filename.split('/')[-1] )
# print(); print()

# schlegel_dust_map_filename = 'lambda_sfd_ebv.fits'
# if not os.path.exists( schlegel_dust_map_filename ):
#     print('WARNING: Cannot find Schlegel et al. dust map with file name:')
#     print(' '*8, schlegel_dust_map_filename)
#     print()
#     print("This can be downloaded via NASA-Goddard's LAMBDA data archive:")
#     print(' '*8, 'http://lambda.gsfc.nasa.gov/product/foreground/f_products.cfm')
#     print('(Look under data > foreground > products > Reddening (E(B-V)) Map,')
#     print('and be sure to download the healpix version.')
#     print()
# else :
#     print('reading Schlegel et al. dust map from file:')
#     print(' '*8, schlegel_dust_map_filename)
#     schlegelDustMap = hp.read_map( schlegel_dust_map_filename, field=0 )
#     hp.mollview( planckDustMap, min=0., max=1., fig=1,
#                  title='Schlegel, Finkbinder & Davis (1998) dust map: %s'
#                  % schlegel_dust_map_filename.split('/')[-1] )
# print(); print()

def healpixAngularCoords( ra, dec ):
    pos = co.SkyCoord( ra*u.deg, dec*u.deg ).galactic
    theta, phi = pi/2. - pos.b.rad, pos.l.rad
    return theta, phi

# def Planck_EBV( theta, phi ):
#     return hp.get_interp_val( planckDustMap, theta, phi )

# def Schlegel_EBV( theta, phi ):
#     return hp.get_interp_val( schlegelDustMap, theta, phi )

def EBV(name, theta, phi):
    """
    Return E(B-V) for given map at given location.

    Valid names are 'planck', 'sfd98' or 'schlafly'.
    """
    success = load_map(name)
    if success:
        return hp.get_interp_val(MAPS[name], theta, phi)
    else:
        return None

# def foregroundCorrection( ra, dec, wavelength ):
#     print('Looking up MW dust redding at (RA, Dec) = (%10.6f, %+10.6f).')
#     theta, phi = healpixAngularCoords( ra, dec )

#     EBV1 = Schlegel_EBV( theta, phi )
#     print('E(B-V) from Schlegel dust map is: %.4f' % EBV1)
#     EBV2 = Planck_EBV( theta, phi )
#     print('E(B-V) from Planck dust map is:   %.4f' % EBV2)

#     correction = MilkyWayDustCorrection( wavelength, EBV2, dustlaw='CCM89' )
#     # this is the multiplicative scaling to correct for foreground dust
#     return correction, EBV1, EBV2

def dustCorrectSAMICube( path, overwrite=False ):
    if not HEALPY_AVAILABLE:
        print('healpy not installed; cannot process dust data.')
        return
    hdulist = pf.open(path, 'update')
    try:
        hdu = hdulist['DUST']
    except KeyError:
        # HDU does not exist; make it
        hdu = pf.ImageHDU()
        hdu.name = 'DUST'
        hdulist.append(hdu)
    else:
        # HDU does exist. Do we want to overwrite it?
        if not overwrite:
            # Don't overwrite; get out of here!
            hdulist.close()
            return
    print('Recording dust data for ' + os.path.basename(path))
    header = hdulist[0].header
    ra, dec = header[ 'CATARA' ], header[ 'CATADEC' ]
    wl = header[ 'CRVAL3' ] + ( header[ 'CDELT3' ] *
                                (1 + np.arange( header[ 'NAXIS3' ] )
                                - header[ 'CRPIX3' ] ))
    theta, phi = healpixAngularCoords( ra, dec )
    for name, map_info in MAPS_FILES.items():
        ebv = EBV(name, theta, phi)
        if ebv is not None:
            # We were able to find the dust map
            hdu.header[map_info['header_key']] = (
                ebv, 'MW reddening E(B-V) from {}'.format(
                    map_info['comment_name']))
            if name == 'planck':
                correction = MilkyWayDustCorrection(wl, ebv)
                hdu.data = correction
        else:
            # We were not able to find the dust map
            print('Warning: {} dust map not available'.format(
                map_info['comment_name']))
            if name == 'planck':
                print('No dust curve recorded')
    hdulist.flush()
    hdulist.close()
    return

def MilkyWayDustCorrection( wavelength, EBV, dustlaw='CCM89' ):
    # MW dust extinction law taken from Cardelli, Clayton & Mathis (1989)
    # my implementation of this follows Madusha's

    # C89 parameterise dust extinction ito a parameter R_v=A_v/E(B-V)
    # here i assume R_v = 3.1; this is Calzetti (2001)'s value for SFers
    # this is also the value given by C89 for the diffuse ISM

    Rv = 3.1

    # C89 give k(lam) normalized to 1 Av = 1. * R_v / E(B-V)
    # i do things in E(B-V) = Av/Rv

    # everything is parameterized according to x = 1 / lambda [um]
    # i assume <wavelength> is given in [Angstroms] = 10000 * [um]
    x = 1./(wavelength/1e4)

    infrared = ( ( 0.3 <= x ) & ( x <= 1.1 ) )
    optical =  ( ( 1.1 <  x ) & ( x <= 3.3 ) )

    a = np.where( infrared, +0.574 * x**1.61, 0. )
    b = np.where( infrared, -0.527 * x**1.61, 0. )

    y = x - 1.82

    # NB. np.polyval does p[0]*X**(N-1) + p[1]*X**(N-2) + ... + p[-2]*X + p[-1]
    # ie. polynmoial coefficients are for higher through to lower powers
    if dustlaw == 'CCM89' :
        acoeffs = ( +0.32999, -0.77530, +0.01979, +0.72085,
                    -0.02427, -0.50447, +0.17699, 1. )
        bcoeffs = ( -2.09002, +5.30260, -0.62251, -5.38434,
                    +1.07233, +2.28305, +1.41338, 0. )
    elif dustlaw == 'OD94' :
        acoeffs = ( -0.505, +1.647, -0.827, -1.718,
                    +1.137, +0.701, -0.609, +0.104, 1. )
        bcoeffs = ( +3.347,-10.805, +5.491,+11.102,
                    -7.985, -3.989, +2.908, +1.952, 0. )
    else :
        print('Do not recognise the given dust law:', dustlaw)
        print('Recognised options are:')
        print('--- CCM89 (Cardelli, Clayton & Mathis, 1989, ApJ 345, 245)')
        print("--- OD94  (O'Donnell, 1994, ApJ 422, 1580")
        print('No dust correction will be performed.')
        return 1.

    a = np.where( optical, np.polyval( acoeffs, y ), a )
    b = np.where( optical, np.polyval( bcoeffs, y ), b )

    attenuation = a + b / Rv
    # Rv is Av / E(B-V) ; correction is normalised to Av = 1
    # so to scale to E(B-V)=1, scale correction by Rv
    attenuation *= Rv

    transmission = 10**( -0.4 * EBV * attenuation )
    # this is the fraction of light transmitted through the foreground
    correction = 1./transmission
    # this is the multiplicative scaling to correct for foreground dust

    return correction

def gamaTest( ):
    import atpy
    print('requires InputCatA.fits; download from GAMA DR2 webpages.')
    gama = atpy.Table( '/Users/ent/data/gama/dr2/InputCatA.fits' )
    ebv0 = gama.EXTINCTION_R / 2.751

    plt.figure( 5 ) ; plt.clf()
    plt.xlabel( 'EBV from GAMA catalogues' )
    plt.ylabel( 'EBV from this code' )
    plt.title( 'red = Schlegel+98 dust map; black = Planck dust map' )

    for i, ( ra, dec ) in enumerate( zip( gama.RA, gama.DEC ) ):
        theta, phi = thetaPhiFromRaDec( ra, dec )
        dust = Planck_EBV( theta, phi )
        dust2 = Schlegel_EBV( theta, phi )

        plt.scatter( gama.EXTINCTION_R[ i ]/3.1, dust, 1, 'k', edgecolors='none' )
        plt.scatter( gama.EXTINCTION_R[ i ]/3.1, dust2, 1, 'r', edgecolors='none' )
        if i % 1000 == 999 :
            plt.draw()

max_stars_repo_name: rscalzo/sami

max_issues_repo_name: rscalzo/sami

max_forks_repo_name: rscalzo/sami

reason: import numpy,import astropy,from astropy


text:
"""A module for TurbidityCurrent2D to produce a grid object from a geotiff
   file or from scratch.

   codeauthor: : Hajime Naruse
"""

from landlab import RasterModelGrid
import numpy as np
from osgeo import gdal, gdalconst
from scipy.ndimage import median_filter
from landlab import FieldError

def create_topography(
    length=8000,
    width=2000,
    spacing=20,
    slope_outside=0.1,
    slope_inside=0.05,
    slope_basin=0.02,
    slope_basin_break=2000,
    canyon_basin_break=2200,
    canyon_center=1000,
    canyon_half_width=100,
    canyon="parabola",
    noise=0.01,
):
    """create an artificial topography where a turbidity current flow down
       A slope and a flat basn plain are set in calculation domain, and a
       parabola or v-shaped canyon is created in the slope.

       Parameters
       ------------------
        length: float, optional
           length of calculation domain [m]

        width: float, optional
           width of calculation domain [m]

        spacing: float, optional
           grid spacing [m]

        slope_outside: float, optional
           topographic inclination in the region outside the canyon

        slope_inside: float, optional
           topographic inclination in the region inside the thalweg of
           the canyon

        slope_basin: float, optional
           topographic inclination of the basin plain

        slope_basin_break: float, optional
           location of slope-basin break

        canyon_basin_break: float, optional
           location of canyon-basin break. This value must be
           larger than slope-basin break.

        canyon_center: float, optional
           location of center of the canyon

        canyon_half_width: float, optional
           half width of the canyon

        canyon: String, optional
           Style of the canyon. 'parabola' or 'V' can be chosen.

        random: float, optional
           Range of random noise to be added on generated topography

        Return
        -------------------------
        grid: RasterModelGrid
           a landlab grid object. Topographic elevation is stored as
           grid.at_node['topographic__elevation']

    """
    # making grid
    # size of calculation domain is 4 x 8 km with dx = 20 m
    lgrids = int(length / spacing)
    wgrids = int(width / spacing)
    grid = RasterModelGrid((lgrids, wgrids), xy_spacing=[spacing, spacing])
    grid.add_zeros("flow__depth", at="node")
    grid.add_zeros("topographic__elevation", at="node")
    grid.add_zeros("flow__horizontal_velocity_at_node", at="node")
    grid.add_zeros("flow__vertical_velocity_at_node", at="node")
    grid.add_zeros("flow__horizontal_velocity", at="link")
    grid.add_zeros("flow__vertical_velocity", at="link")
    grid.add_zeros("bed__thickness", at="node")

    # making topography
    # set the slope
    grid.at_node["topographic__elevation"] = (
        grid.node_y - slope_basin_break
    ) * slope_outside

    if canyon == "parabola":
        # set canyon
        d0 = slope_inside * (canyon_basin_break - slope_basin_break)
        d = slope_inside * (grid.node_y - canyon_basin_break) - d0
        a = d0 / canyon_half_width ** 2
        canyon_elev = a * (grid.node_x - canyon_center) ** 2 + d
        inside = np.where(canyon_elev < grid.at_node["topographic__elevation"])
        grid.at_node["topographic__elevation"][inside] = canyon_elev[inside]

    # set basin
    basin_height = (grid.node_y - slope_basin_break) * slope_basin
    basin_region = grid.at_node["topographic__elevation"] < basin_height
    grid.at_node["topographic__elevation"][basin_region] = basin_height[basin_region]

    # add random value on topographic elevation (+- noise)
    grid.at_node["topographic__elevation"] += (
        2.0 * noise * (np.random.rand(grid.number_of_nodes) - 0.5)
    )

    grid.set_closed_boundaries_at_grid_edges(False, False, False, False)

    return grid

def create_init_flow_region(
    grid,
    initial_flow_concentration=0.02,
    initial_flow_thickness=200,
    initial_region_radius=200,
    initial_region_center=[1000, 7000],
):
    """ making initial flow region in a grid, assuming lock-exchange type initiation
         of a turbidity current. Plan-view morphology of a suspended cloud is a circle,

         Parameters
         ----------------------
         grid: RasterModelGrid
            a landlab grid object

         initial_flow_concentration: float, optional
            initial flow concentration

         initial_flow_thickness: float, optional
            initial flow thickness

         initial_region_radius: float, optional
            radius of initial flow region

         initial_region_center: list, optional
            [x, y] coordinates of center of initial flow region
    """
    # check number of grain size classes
    if type(initial_flow_concentration) is float:
        initial_flow_concentration_i = np.array([initial_flow_concentration])
    else:
        initial_flow_concentration_i = np.array(initial_flow_concentration).reshape(
            len(initial_flow_concentration), 1
        )

    # initialize flow parameters
    for i in range(len(initial_flow_concentration_i)):
        try:
            grid.add_zeros("flow__sediment_concentration_{}".format(i), at="node")
        except FieldError:
            grid.at_node["flow__sediment_concentration_{}".format(i)][:] = 0.0
        try:
            grid.add_zeros("bed__sediment_volume_per_unit_area_{}".format(i), at="node")
        except FieldError:
            grid.at_node["bed__sediment_volume_per_unit_area_{}".format(i)][:] = 0.0

    try:
        grid.add_zeros("flow__sediment_concentration_total", at="node")
    except FieldError:
        grid.at_node["flow__sediment_concentration_total"][:] = 0.0
    try:
        grid.add_zeros("flow__depth", at="node")
    except FieldError:
        grid.at_node["flow__depth"][:] = 0.0
    try:
        grid.add_zeros("flow__horizontal_velocity_at_node", at="node")
    except FieldError:
        grid.at_node["flow__horizontal_velocity_at_node"][:] = 0.0
    try:
        grid.add_zeros("flow__vertical_velocity_at_node", at="node")
    except FieldError:
        grid.at_node["flow__vertical_velocity_at_node"][:] = 0.0
    try:
        grid.add_zeros("flow__horizontal_velocity", at="link")
    except FieldError:
        grid.at_link["flow__horizontal_velocity"][:] = 0.0
    try:
        grid.add_zeros("flow__vertical_velocity", at="link")
    except FieldError:
        grid.at_link["flow__vertical_velocity"][:] = 0.0

    # set initial flow region
    initial_flow_region = (
        (grid.node_x - initial_region_center[0]) ** 2
        + (grid.node_y - initial_region_center[1]) ** 2
    ) < initial_region_radius ** 2
    grid.at_node["flow__depth"][initial_flow_region] = initial_flow_thickness
    grid.at_node["flow__depth"][~initial_flow_region] = 0.0
    for i in range(len(initial_flow_concentration_i)):
        grid.at_node["flow__sediment_concentration_{}".format(i)][
            initial_flow_region
        ] = initial_flow_concentration_i[i]
        grid.at_node["flow__sediment_concentration_{}".format(i)][
            ~initial_flow_region
        ] = 0.0
    grid.at_node["flow__sediment_concentration_total"][initial_flow_region] = np.sum(
        initial_flow_concentration_i
    )

def create_topography_from_geotiff(
    geotiff_filename, xlim=None, ylim=None, spacing=500, filter_size=[1, 1]
):
    """create a landlab grid file from a geotiff file

       Parameters
       -----------------------
       geotiff_filename: String
          name of a geotiff-format file to import

       xlim: list, optional
          list [xmin, xmax] to specify x coordinates of a region of interest
             in a geotiff file to import

       ylim: list, optional
          list [ymin, ymax] to specify y coordinates of a region of interest
             in a geotiff file to import

       spacing: float, optional
          grid spacing

       filter_size: list, optional
          [x, y] size of a window used in a median filter.
            This filter is applied for smoothing DEM data.

       Return
       ------------------------
       grid: RasterModelGrid
          a landlab grid object to be used in TurbidityCurrent2D

    """

    # read a geotiff file into ndarray
    topo_file = gdal.Open(geotiff_filename, gdalconst.GA_ReadOnly)
    topo_data = topo_file.GetRasterBand(1).ReadAsArray()
    if (xlim is not None) and (ylim is not None):
        topo_data = topo_data[xlim[0] : xlim[1], ylim[0] : ylim[1]]

    # Smoothing by median filter
    topo_data = median_filter(topo_data, size=filter_size)

    grid = RasterModelGrid(topo_data.shape, xy_spacing=[spacing, spacing])
    grid.add_zeros("flow__depth", at="node")
    grid.add_zeros("topographic__elevation", at="node")
    grid.add_zeros("flow__horizontal_velocity", at="link")
    grid.add_zeros("flow__vertical_velocity", at="link")
    grid.add_zeros("bed__thickness", at="node")
    grid.at_node["topographic__elevation"][grid.nodes] = topo_data

    return grid


text:
import time
import copy
import random
import json
import argparse
import datetime
import numpy as np
import torch
from torch.optim import AdamW
from torch.utils.data import DataLoader
from tensorboardX import SummaryWriter
from sklearn.model_selection import train_test_split
import pandas as pd
from tqdm import tqdm
import torch.nn as nn
from sklearn import preprocessing

from tensorboardX import SummaryWriter
from torch.cuda.amp import GradScaler, autocast

from types_ import *
from utils import *
from losses import *
from read_data import *

class RNAEncoder(nn.Module):
    def __init__(self,
                 in_channels: int,
                 hidden_dims: List):
        super(RNAEncoder, self).__init__()

        self.in_channels = in_channels

        modules = []
        # Build encoder
        for h_dim in hidden_dims:
            modules.append(
                nn.Sequential(
                    nn.Linear(in_channels, h_dim),
                    nn.LeakyReLU()
                    )
            )
            in_channels = h_dim

        self.encoder = nn.Sequential(*modules)

    def forward(self, x):
        return self.encoder(x)

class RNAModel(nn.Module):
    def __init__(self,
                 in_channels: int,
                 hidden_dims: List,
                 out_channels: int):
        super(RNAModel, self).__init__()

        self.encoder = RNAEncoder(in_channels, hidden_dims)
        self.fc = nn.Sequential(
            nn.Linear(hidden_dims[-1], 200),
            nn.LeakyReLU(),
            nn.Linear(200, out_channels))

    def forward(self, x):
        x = self.encoder(x)
        return self.fc(x)

def train_RNA(model, criterion, optimizer, dataloaders,
          save_dir='checkpoints/models/', device=None,
          log_interval=100, summary_writer=None, num_epochs=100,
          verbose=True):

     since = time.time()

     best_model_wts = copy.deepcopy(model.state_dict())
     best_epoch = 0
     best_loss = 0
     loss_array = {'train': [], 'val': []}

     global_summary_step = {'train': 0, 'val': 0}
     scaler = GradScaler()
     for epoch in range(num_epochs):
         print('Epoch {}/{}'.format(epoch, num_epochs - 1))
         print('-' * 10)

         sizes = {'train': 0, 'val': 0}
         inputs_seen = {'train': 0, 'val': 0}

         for phase in ['train', 'val']:
            if phase == 'train':
                 model.train()
            else:
                 model.eval()

            running_loss = 0.0
            running_corrects = 0.0
            summary_step = global_summary_step[phase]
            # for logging tensorboard
            last_running_loss = 0.0
            output_list = []
            vital_status_list = []
            survival_months_list = []
            for b_idx, batch in tqdm(enumerate(dataloaders[phase])):
                vital_status_list.append(batch[1].numpy().flatten())
                survival_months_list.append(batch[2].numpy().flatten())
                vital_status = batch[1].to(device)
                survival_months = batch[2].to(device)

                rna_data = batch[0].to(device)
                optimizer.zero_grad()
                with torch.set_grad_enabled(phase=='train'):
                    outputs = model(rna_data)
                    output_list.append(outputs.detach().cpu().numpy().flatten())
                    loss = criterion(outputs, vital_status.view(vital_status.size(0)), survival_months.view(survival_months.size(0)))
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                summary_step += 1
                running_loss += loss.item()
                sizes[phase] += rna_data.size(0)
                inputs_seen[phase] += rna_data.size(0)

                # Emptying memory
                del outputs
                torch.cuda.empty_cache()

                if (summary_step % log_interval == 0):
                    loss_to_log = (running_loss - last_running_loss) / inputs_seen[phase]

                    if summary_writer is not None:
                        summary_writer.add_scalar("{}/loss".format(phase), loss_to_log, summary_step)

                    last_running_loss = running_loss
                    last_running_corrects = running_corrects
                    inputs_seen[phase] = 0.0

            global_summary_step[phase] = summary_step
            epoch_loss = running_loss / sizes[phase]

            loss_array[phase].append(epoch_loss)
            output_list = np.concatenate(output_list, axis=0)

            survival_months_list = np.concatenate(survival_months_list, axis=0)
            vital_status_list = np.concatenate(vital_status_list, axis=0)
            CI = get_survival_CI(output_list, survival_months_list, vital_status_list)
            if verbose:
                print('{} Loss: {:.4f}, CI: {:.4f}'.format(
                    phase, epoch_loss, CI))

            if phase == 'val' and CI > best_loss:
                best_loss = epoch_loss
                torch.save(model.state_dict(), os.path.join(save_dir, 'model_dict_best.pt'))
                best_epoch = epoch

     torch.save(model.state_dict(), os.path.join(save_dir, 'model_last.pt'))

     time_elapsed = time.time() - since

     model.load_state_dict(torch.load(os.path.join(save_dir, 'model_dict_best.pt')))

     results = {
         'best_epoch': best_epoch,
         'best_loss': best_loss
     }
     return model, results

def evaluate_RNA(model, dataloader, dataset_size, criterion, device=None):
    model.eval()

    running_loss = 0
    running_acc = 0
    output_list = []
    vital_status_list = []
    survival_months_list = []
    for batch in tqdm(dataloader):
        vital_status_list.append(batch[1].numpy().flatten())
        survival_months_list.append(batch[2].numpy().flatten())
        vital_status = batch[1].to(device)
        survival_months = batch[2].to(device)

        rna_data = batch[0].to(device)
        with torch.set_grad_enabled(False):
            outputs = model(rna_data)
            loss = criterion(outputs, vital_status.view(vital_status.size(0)), survival_months.view(survival_months.size(0)))
            running_loss += loss.item()
            output_list.append(outputs.detach().cpu().numpy().flatten())

    test_loss = running_loss / dataset_size

    output_list = np.concatenate(output_list, axis=0)
    survival_months_list = np.concatenate(survival_months_list, axis=0)
    vital_status_list = np.concatenate(vital_status_list, axis=0)

    CI = get_survival_CI(output_list, survival_months_list, vital_status_list)
    print('Loss {}; CI {}'.format(test_loss, CI))
    test_results = {
        'preds': output_list,
        'survival_months': survival_months_list,
        'vital_status': vital_status_list,
        'CI': CI,
        'loss': test_loss
    }

    return test_results

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='RNA-Seq Training')
    parser.add_argument('--config', type=str, help='JSON config file')
    parser.add_argument('--checkpoint', type=str, default=None,
            help='File with the checkpoint to start with')
    parser.add_argument('--save_dir', type=str, default=None,
            help='Where to save the checkpoints')
    parser.add_argument('--flag', type=str, default=None,
            help='Flag to use for saving the checkpoints')
    parser.add_argument('--seed', type=int, default=99,
            help='Seed for random generation')
    parser.add_argument('--log', type=int, default=0,
            help='Use tensorboard for experiment logging')
    parser.add_argument('--parallel', type=int, default=0,
            help='Use DataParallel training')
    parser.add_argument('--fp16', type=int, default=0,
            help='Use mixed-precision training')
    parser.add_argument('--batch_size', type=int, default=16,
                    help='Batch size')
    parser.add_argument('--lr', type=float, default=3e-3,
                        help='Learning rate to use')
    parser.add_argument('--quick', type=int, default=0,
            help='If a subset of samples is used for a quick training')
    parser.add_argument('--num_epochs', type=int, default=100,
            help='Number of epochs to train the model')
    parser.add_argument('--num_rna_features', type=int, default=17655,
            help='Number of RNA features used as input')
    parser.add_argument('--num_samples', type=int, default=20,
            help='Number of samples to use to train the model')
    parser.add_argument('--problem', type=str, default='survival',
            help='Problem to classify')
    args = parser.parse_args()

    np.random.seed(args.seed)
    torch.manual_seed(args.seed)

    with open(args.config) as f:
        config = json.load(f)

    print(10*'-')
    print('Config for this experiment \n')
    print(config)
    print(args)
    print(10*'-')

    if 'flag' in config:
        args.flag = config['flag']
    else:
        args.flag = 'train_{date:%Y-%m-%d %H:%M:%S}'.format(date=datetime.datetime.now())

    if not os.path.exists(args.save_dir):
        os.mkdir(args.save_dir)

    path_csv = config['path_csv']
    rna_features = args.num_rna_features
    quick = args.quick
    batch_size = args.batch_size
    print('Loading dataset...')

    df = pd.read_csv(path_csv)

    test_CIs = []
    for i in range(10):
        try:
            train_df, test_df = train_test_split(df, test_size=0.2, stratify=df['tcga_project'], random_state=args.seed)

            train_df, val_df = train_test_split(train_df, test_size=0.2, stratify=train_df['tcga_project'], random_state=args.seed)

            train_df, val_df, test_df, scaler = normalize_dfs(train_df, val_df, test_df)

            train_dataset = RNADataset(train_df, quick=quick, num_samples=args.num_samples, seed=args.seed*i)

            val_dataset = RNADataset(val_df, quick=0, num_samples=64)

            test_dataset = RNADataset(test_df, quick=0, num_samples=64)

            train_dataloader = DataLoader(train_dataset, batch_size=batch_size,
                        num_workers=config['n_workers'], pin_memory=True, shuffle=True)
            val_dataloader = DataLoader(val_dataset, batch_size=batch_size, num_workers=config['n_workers'],
            pin_memory=True, shuffle=False)
            test_dataloader = DataLoader(test_dataset, batch_size=batch_size,
            num_workers=config['n_workers'], pin_memory=True, shuffle=False)

            dataloaders = {
                    'train': train_dataloader,
                    'val': val_dataloader}

            dataset_sizes = {
                    'train': len(train_dataset),
                    'val': len(val_dataset)
                    }

            print('Finished loading dataset and creating dataloader')

            print('Initializing models')

            if args.problem == 'survival':
                model = RNAModel(in_channels=rna_features,
                                hidden_dims=[6000,512],
                                out_channels=1)
            else:
                model = RNAModel(in_channels=rna_features,
                                hidden_dims=[6000,512],
                                out_channels=len(classes))
            if args.checkpoint is not None:
                model.fc.apply(init_weights_xavier)
            else:
                model.apply(init_weights_xavier)

            if args.checkpoint is not None:
                print('Restoring from checkpoint')
                print(args.checkpoint)
                model.encoder.load_state_dict(torch.load(args.checkpoint))
                print('Loaded model from checkpoint')

            model = model.to(config['device'])
            # add optimizer
            lr = args.lr
            optimizer = AdamW(model.parameters(), weight_decay = 0.1, lr=lr)

            # add loss function
            if args.problem == 'survival':
                criterion = CoxLoss()
            else:
                criterion = nn.CrossEntropyLoss()

            # train model

            if args.log:
                summary_writer = SummaryWriter(
                        os.path.join(config['summary_path'],
                            datetime.datetime.now().strftime("%Y-%m-%d_%H:%M:%S") + "_{0}".format(args.flag)))

                summary_writer.add_text('config', str(config))
            else:
                summary_writer = None

            model, trainval_results = train_RNA(model, criterion, optimizer, dataloaders,
                                                save_dir=args.save_dir, device='cuda:0',
                                                num_epochs=args.num_epochs)

            test_results = evaluate_RNA(model, test_dataloader, len(test_dataset),
                                        criterion, device='cuda:0')

            test_CIs.append(test_results['CI'])
            with open(args.save_dir+'test_predictions.pkl', 'wb') as f:
                pickle.dump(test_results, f)
        except:
            continue

    print('Mean {}; std {}'.format(np.mean(test_CIs), np.std(test_CIs)))
    print(test_CIs)

max_stars_repo_name: pacocp/SSL-Survival

max_issues_repo_name: pacocp/SSL-Survival

max_forks_repo_name: pacocp/SSL-Survival