<commit_before>#!/usr/bin/env python

import numpy as np
from scipy.interpolate import interp1d
from scipy.integrate import simps
from scipy.signal import convolve
from scipy.stats import norm,lognorm
from astropy.table import Table

# XXX
from astropy.io.fits import getdata
from astropy.io import ascii as ascii_io

from .sqbase import datadir,mag2lum
from . import dustextinction


##############################################################################
# Samplers
##############################################################################

class Sampler(object):
	def __init__(self,low,high):
		self.low = low
		self.high = high
	def sample(self,n):
		raise NotImplementedError
	def resample(self,qsoData,**kwargs):
		pass
	def __call__(self,n):
		return self.sample(n)
	def __str__(self):
		s = str((self.low,self.high))
		return s

class FixedSampler(Sampler):
	def __init__(self,vals):
		self.low = None
		self.high = None
		self.vals = vals
	def sample(self,n):
		return self.vals

class ConstSampler(Sampler):
	def __init__(self,val):
		self.low = None
		self.high = None
		self.val = val
	def sample(self,n):
		return np.repeat(self.val,n)

class UniformSampler(Sampler):
	def sample(self,n):
		return np.linspace(self.low,self.high,n)

class GridSampler(Sampler):
	def __init__(self,low,high,nbins=None,stepsize=None):
		if nbins is None and stepsize is None:
			raise ValueError("Must specify nbins or stepsize")
		super(GridSampler,self).__init__(low,high)
		if stepsize:
			nbins = (high - low) / stepsize
		self.nbins = nbins+1
	def sample(self,n):
		arr = np.linspace(self.low,self.high,n*self.nbins)
		return arr.reshape(self.nbins,n)

class CdfSampler(Sampler):
	def _init_cdf(self):
		self.cdf_low = self.rv.cdf(self.low)
		self.cdf_high = self.rv.cdf(self.high)
	def _getpoints(self,x):
		return self.cdf_low + (self.cdf_high-self.cdf_low)*x
	def _sample(self,x):
		return self.rv.ppf(x)
	def sample(self,n):
		x = np.random.random(n)
		return self._sample(self._getpoints(x))

class PowerLawSampler(CdfSampler):
	def __init__(self,low,high,a):
		self.rv = self
		super(PowerLawSampler,self).__init__(low,high)
		self.a = a
	def cdf(self,x):
		x1,x2,a = self.low,self.high,self.a
		if np.any(x<x1) or np.any(x>x2):
			raise ValueError
		return (x**(a+1) - x1**(a+1)) / (a+1)
	def ppf(self,y):
		if np.any(y<0) or np.any(y>1):
			raise ValueError
		x1,x2,a = self.low,self.high,self.a
		return np.power( (x2**(a+1)-x1**(a+1))*y + x1**(a+1), (a+1)**-1 )

class GaussianSampler(CdfSampler):
	def __init__(self,mean,sigma,low=-np.inf,high=np.inf):
		super(GaussianSampler,self).__init__(low,high)
		self.mean = mean
		self.sigma = sigma
		self._reset()
		self._init_cdf()
	def _reset(self):
		self.rv = norm(loc=self.mean,scale=self.sigma)

class LogNormalSampler(CdfSampler):
	def __init__(self,mean,sigma,low=-np.inf,high=np.inf):
		super(LogNormalSampler,self).__init__(low,high)
		self.mean = mean
		self.sigma = sigma
		self.rv = lognorm(loc=self.mean,scale=self.sigma)
		self._init_cdf()

class DoublePowerLawSampler(Sampler):
	def __init__(self,a,b,x0,low=-np.inf,high=np.inf):
		super(DoublePowerLawSampler,self).__init__(low,high)
		self.a = a
		self.b = b
		self.x0 = x0
	def sample(self,n):
		raise NotImplementedError

class LinearTrendWithAsymScatterSampler(Sampler):
	def __init__(self,coeffs,x,low=-np.inf,high=np.inf):
		super(LinearTrendWithAsymScatterSampler,self).__init__(low,high)
		self.coeffs = coeffs
		self._reset(x)
	def _reset(self,x):
		xmn,xlo,xhi = np.polyval(self.coeffs.T,x)
		if xmn-xlo < 1e-10:
			self.loSampler = ConstSampler(xmn)
		else:
			self.loSampler = GaussianSampler(xmn,xmn-xlo,
			                                 low=self.low,high=self.high)
		if xhi-xmn < 1e-10:
			self.hiSampler = ConstSampler(xmn)
		else:
			self.hiSampler = GaussianSampler(xmn,xhi-xmn,
			                                 low=self.low,high=self.high)
	def _getpoints(self,x):
		if isinstance(self.loSampler,ConstSampler):
			xlo = self.loSampler.sample(len(x))
		else:
			xlo = self.loSampler._getpoints(x)
		if isinstance(self.hiSampler,ConstSampler):
			xhi = self.hiSampler.sample(len(x))
		else:
			xhi = self.hiSampler._getpoints(x)
		return np.choose(x<0,[xlo,xhi])

class BaldwinEffectSampler(LinearTrendWithAsymScatterSampler):
	def __init__(self,coeffs,absMag,low=0,high=np.inf):
		super(BaldwinEffectSampler,self).__init__(coeffs,absMag,
		                                          low=low,high=high)
	def sample(self,n):
		# save the x values for reuse
		self.x = np.random.random(n)
		return self._getpoints(self.x)
	def resample(self,qsoData,**kwargs):
		m = qsoData.absMag - Mref
		self._reset(m)
		return self._getpoints(self.x)



##############################################################################
# Simulation variables
##############################################################################

class QsoSimVar(object):
	def __init__(self,sampler,name=None):
		self.sampler = sampler
		if name is not None:
			self.name = name
	def __call__(self,n):
		return self.sampler(n)
	def __str__(self):
		return str(self.sampler)

class AppMagVar(QsoSimVar):
	name = 'appMag'
	def __init__(self,sampler,band=None):
		super(AppMagVar,self).__init__(sampler)
		self.band = band

class AbsMagVar(QsoSimVar):
	name = 'absMag'
	def __init__(self,sampler,restWave=None):
		'''if restWave is none then bolometric'''
		super(AbsMagVar,self).__init__(sampler)<commit_msg>fix problems with beff sampler<commit_after>#!/usr/bin/env python

import numpy as np
from scipy.interpolate import interp1d
from scipy.integrate import simps
from scipy.signal import convolve
from scipy.stats import norm,lognorm
from astropy.table import Table

# XXX
from astropy.io.fits import getdata
from astropy.io import ascii as ascii_io

from .sqbase import datadir,mag2lum
from . import dustextinction


##############################################################################
# Samplers
##############################################################################

class Sampler(object):
	def __init__(self,low,high):
		self.low = low
		self.high = high
	def sample(self,n):
		raise NotImplementedError
	def resample(self,qsoData,**kwargs):
		pass
	def __call__(self,n):
		return self.sample(n)
	def __str__(self):
		s = str((self.low,self.high))
		return s

class FixedSampler(Sampler):
	def __init__(self,vals):
		self.low = None
		self.high = None
		self.vals = vals
	def sample(self,n):
		return self.vals

class ConstSampler(Sampler):
	def __init__(self,val):
		self.low = None
		self.high = None
		self.val = val
	def sample(self,n):
		return np.repeat(self.val,n)

class UniformSampler(Sampler):
	def sample(self,n):
		return np.linspace(self.low,self.high,n)

class GridSampler(Sampler):
	def __init__(self,low,high,nbins=None,stepsize=None):
		if nbins is None and stepsize is None:
			raise ValueError("Must specify nbins or stepsize")
		super(GridSampler,self).__init__(low,high)
		if stepsize:
			nbins = (high - low) / stepsize
		self.nbins = nbins+1
	def sample(self,n):
		arr = np.linspace(self.low,self.high,n*self.nbins)
		return arr.reshape(self.nbins,n)

class CdfSampler(Sampler):
	def _init_cdf(self):
		self.cdf_low = self.rv.cdf(self.low)
		self.cdf_high = self.rv.cdf(self.high)
	def _getpoints(self,x):
		return self.cdf_low + (self.cdf_high-self.cdf_low)*x
	def _sample(self,x):
		return self.rv.ppf(x)
	def sample(self,n):
		x = np.random.random(n)
		return self._sample(self._getpoints(x))

class PowerLawSampler(CdfSampler):
	def __init__(self,low,high,a):
		self.rv = self
		super(PowerLawSampler,self).__init__(low,high)
		self.a = a
	def cdf(self,x):
		x1,x2,a = self.low,self.high,self.a
		if np.any(x<x1) or np.any(x>x2):
			raise ValueError
		return (x**(a+1) - x1**(a+1)) / (a+1)
	def ppf(self,y):
		if np.any(y<0) or np.any(y>1):
			raise ValueError
		x1,x2,a = self.low,self.high,self.a
		return np.power( (x2**(a+1)-x1**(a+1))*y + x1**(a+1), (a+1)**-1 )

class GaussianSampler(CdfSampler):
	def __init__(self,mean,sigma,low=-np.inf,high=np.inf):
		super(GaussianSampler,self).__init__(low,high)
		self.mean = mean
		self.sigma = sigma
		self._reset()
		self._init_cdf()
	def _reset(self):
		self.rv = norm(loc=self.mean,scale=self.sigma)

class LogNormalSampler(CdfSampler):
	def __init__(self,mean,sigma,low=-np.inf,high=np.inf):
		super(LogNormalSampler,self).__init__(low,high)
		self.mean = mean
		self.sigma = sigma
		self.rv = lognorm(loc=self.mean,scale=self.sigma)
		self._init_cdf()

class DoublePowerLawSampler(Sampler):
	def __init__(self,a,b,x0,low=-np.inf,high=np.inf):
		super(DoublePowerLawSampler,self).__init__(low,high)
		self.a = a
		self.b = b
		self.x0 = x0
	def sample(self,n):
		raise NotImplementedError

class LinearTrendWithAsymScatterSampler(Sampler):
	def __init__(self,coeffs,pts,low=-np.inf,high=np.inf):
		super(LinearTrendWithAsymScatterSampler,self).__init__(low,high)
		self.coeffs = coeffs
		self.npts = len(pts)
		self._reset(pts)
	def _reset(self,pts):
		xmn,xlo,xhi = [ np.polyval(c,pts) for c in self.coeffs ]
		siglo = np.clip(xmn-xlo,1e-10,np.inf)
		sighi = np.clip(xhi-xmn,1e-10,np.inf)
		self.loSampler = GaussianSampler(xmn,siglo,
		                                 low=self.low,high=self.high)
		self.hiSampler = GaussianSampler(xmn,sighi,
		                                 low=self.low,high=self.high)
	def _sample(self,x):
		# XXX this doesn't seem to pick correctly
		xlo = self.loSampler._sample(self.loSampler._getpoints(x))
		xhi = self.hiSampler._sample(self.hiSampler._getpoints(x))
		return np.choose(x<0.5,[xlo,xhi])

class BaldwinEffectSampler(LinearTrendWithAsymScatterSampler):
	def __init__(self,coeffs,absMag,low=0,high=np.inf):
		super(BaldwinEffectSampler,self).__init__(coeffs,absMag,
		                                          low=low,high=high)
	def sample(self,n):
		if n != self.npts:
			raise ValueError("BaldwinEffectSampler input does not match "
			                 "preset (%d != %d)" % (n,self.npts))
		# save the x values for reuse
		self.x = np.random.random(n)
		return self._sample(self.x)
	def resample(self,absMag,**kwargs):
		self._reset(absMag)
		return self._sample(self.x)



##############################################################################
# Simulation variables
##############################################################################

class QsoSimVar(object):
	def __init__(self,sampler,name=None):
		self.sampler = sampler
		if name is not None:
			self.name = name
	def __call__(self,n):
		return self.sampler(n)
	def __str__(self):
		return str(self.sampler)

class AppMagVar(QsoSimVar):
	name = 'appMag'
	def __init__(self,sampler,band=None):
		super(AppMagVar,self).__init__(sampler)
		self.band = band

class AbsMagVar(QsoSimVar):
	name = 'absMag'
	def __init__(self,sampler,restWave=None):
		'''if restWave is none then bolometric'''
		super(AbsMagVar,self).__init__(sampler)
{'fraction_non_alphanumeric': 0.18763208674078838, 'fraction_numerical': 0.006340163557842507, 'mean_word_length': 4.66875, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '185467', 'n_tokens_mistral': 4395, 'n_tokens_neox': 4092, 'n_words': 726}