Commit 87535b99 authored by Philipp Arras's avatar Philipp Arras

Move stuff from global newton to nifty

parent 77f85939
......@@ -7,6 +7,8 @@ from .hp_space import HPSpace
from .gl_space import GLSpace
from .dof_space import DOFSpace
from .power_space import PowerSpace
from .log_rg_space import LogRGSpace
__all__ = ["Domain", "UnstructuredDomain", "StructuredDomain", "RGSpace",
"LMSpace", "HPSpace", "GLSpace", "DOFSpace", "PowerSpace"]
"LMSpace", "HPSpace", "GLSpace", "DOFSpace", "PowerSpace",
"LogRGSpace"]
from functools import reduce
from ..sugar import exp
import numpy as np
from ..dobj import ibegin
from ..field import Field
from .structured_domain import StructuredDomain
class LogRGSpace(StructuredDomain):
_needed_for_hash = ['_shape', '_bindistances', '_t_0', '_harmonic']
def __init__(self, shape, bindistances, t_0, harmonic=False):
super(LogRGSpace, self).__init__()
self._harmonic = bool(harmonic)
if np.isscalar(shape):
shape = (shape,)
self._shape = tuple(int(i) for i in shape)
self._bindistances = tuple(bindistances)
self._t_0 = tuple(t_0)
self._dim = int(reduce(lambda x, y: x * y, self._shape))
self._dvol = float(reduce(lambda x, y: x * y, self._bindistances))
@property
def harmonic(self):
return self._harmonic
@property
def shape(self):
return self._shape
def scalar_dvol(self):
return self._dvol
@property
def bindistances(self):
return np.array(self._bindistances)
@property
def size(self):
return np.prod(self._shape)
@property
def t_0(self):
return np.array(self._t_0)
def __repr__(self):
return ("LogRGSpace(shape=%r, harmonic=%r)"
% (self.shape, self.harmonic))
def get_default_codomain(self):
if self._harmonic:
raise ValueError("only supported for nonharmonic space")
codomain_bindistances = 1. / (self.bindistances * self.shape)
return LogRGSpace(self.shape, codomain_bindistances,
np.zeros(len(self.shape)), True)
def get_k_length_array(self):
out = Field(self, dtype=np.float64)
oloc = out.local_data
ib = ibegin(out.val)
res = np.arange(oloc.shape[0], dtype=np.float64) + ib[0]
res = np.minimum(res, self.shape[0]-res)*self.bindistances[0]
if len(self.shape) == 1:
oloc[()] = res
return out
res *= res
for i in range(1, len(self.shape)):
tmp = np.arange(oloc.shape[i], dtype=np.float64) + ib[i]
tmp = np.minimum(tmp, self.shape[i]-tmp)*self.bindistances[i]
tmp *= tmp
res = np.add.outer(res, tmp)
oloc[()] = np.sqrt(res)
return out
def get_expk_length_array(self):
# FIXME This is a hack! Only for plotting. Seems not to be the final version.
out = exp(self.get_k_length_array())
out.val[1:] = out.val[:-1]
out.val[0] = 0
return out
from .amplitude_model import make_amplitude_model
from .apply_data import ApplyData
from .los_response import LOSResponse
from .noise_energy import NoiseEnergy
from .nonlinear_power_energy import NonlinearPowerEnergy
from .nonlinear_wiener_filter_energy import NonlinearWienerFilterEnergy
from .nonlinearities import Exponential, Linear, PositiveTanh, Tanh
from .poisson_energy import PoissonEnergy
from .unit_log_gauss import UnitLogGauss
from .poisson_log_likelihood import PoissonLogLikelihood
from .smooth_sky import make_smooth_mf_sky_model, make_smooth_sky_model
from .unit_log_gauss import UnitLogGauss
from .wiener_filter_curvature import WienerFilterCurvature
from .wiener_filter_energy import WienerFilterEnergy
import numpy as np
from ..domains import PowerSpace, UnstructuredDomain
from ..field import Field
from ..multi import MultiField
from ..sugar import makeOp, sqrt
def _ceps_kernel(dof_space, k, a, k0):
return a**2/(1+(k/(k0*dof_space.bindistances[0]))**2)**2
def make_amplitude_model(s_space, Npixdof, ceps_a, ceps_k, sm, sv, im, iv,
keys=['tau', 'phi']):
'''
Method for construction of amplitude model
Computes a smooth power spectrum.
Output lives in PowerSpace.
Parameters
----------
Npixdof : #pix in dof_space
ceps_a, ceps_k0 : Smoothnessparameters in ceps_kernel
eg. ceps_kernel(k) = (a/(1+(k/k0)**2))**2
a = ceps_a, k0 = ceps_k0
sm, sv : slope_mean = expected exponent of powerlaw (e.g. -4),
slope_variance (default=1)
im, iv : y-intercept_mean, y-intercept_variance of power_slope
'''
from ..operators import (ExpTransform, QHTOperator, SlopeOperator,
SymmetrizingOperator)
h_space = s_space.get_default_codomain()
p_space = PowerSpace(h_space)
exp_transform = ExpTransform(p_space, Npixdof)
logk_space = exp_transform.domain[0]
dof_space = logk_space.get_default_codomain()
param_space = UnstructuredDomain(2)
qht = QHTOperator(dof_space, logk_space)
sym = SymmetrizingOperator(logk_space)
phi_mean = np.array([sm, im])
phi_sig = np.array([sv, iv])
slope = SlopeOperator(param_space, logk_space, phi_sig)
norm_phi_mean = Field(param_space, val=phi_mean/phi_sig)
fields = {keys[0]: Field.from_random('normal', dof_space),
keys[1]: Field.from_random('normal', param_space)}
position = MultiField(fields)
dof_space = position[keys[0]].domain[0]
kern = lambda k: _ceps_kernel(dof_space, k, ceps_a, ceps_k)
cepstrum = create_cepstrum_amplitude_field(dof_space, kern)
ceps = makeOp(sqrt(cepstrum))
smooth_op = sym * qht * ceps
from ..models import Variable
smooth_spec = smooth_op(Variable(position)[keys[0]])
phi = Variable(position)[keys[1]] + norm_phi_mean
linear_spec = slope(phi)
loglog_spec = smooth_spec + linear_spec
from ..models import PointwiseExponential
xlog_ampl = PointwiseExponential(0.5*loglog_spec)
internals = {'loglog_spec': loglog_spec,
'qht': qht,
'ceps': ceps,
'norm_phi_mean': norm_phi_mean}
return exp_transform(xlog_ampl), internals
def create_cepstrum_amplitude_field(domain, cepstrum):
"""Creates a ...
Writes the sum of all modes into the zero-mode.
Parameters
----------
domain: ???
???
cepstrum: Callable
???
"""
dim = len(domain.shape)
dist = domain.bindistances
shape = domain.shape
# Prepare q_array
q_array = np.zeros((dim,) + shape)
if dim == 1:
ks = domain.get_k_length_array().val
q_array = np.array([ks])
else:
for i in range(dim):
ks = np.minimum(shape[i] - np.arange(shape[i]) +
1, np.arange(shape[i])) * dist[i]
fst_dims = (1,) * i
lst_dims = (1,) * (dim - i - 1)
q_array[i] += ks.reshape(fst_dims + (shape[i],) + lst_dims)
# Fill cepstrum field (all non-zero modes)
no_zero_modes = (slice(1, None),) * dim
ks = q_array[(slice(None),) + no_zero_modes]
cepstrum_field = np.zeros(shape)
cepstrum_field[no_zero_modes] = cepstrum(ks)
# Fill cepstrum field (zero-mode subspaces)
for i in range(dim):
# Prepare indices
fst_dims = (slice(None),) * i
lst_dims = (slice(None),) * (dim - i - 1)
sl = fst_dims + (slice(1, None),) + lst_dims
sl2 = fst_dims + (0,) + lst_dims
# Do summation
cepstrum_field[sl2] = np.sum(cepstrum_field[sl], axis=i)
return Field(domain, val=cepstrum_field)
def ApplyData(data, var, model_data):
from .. import DiagonalOperator, Constant, sqrt
sqrt_n = DiagonalOperator(sqrt(var))
data = Constant(model_data.position, data)
return sqrt_n.inverse(model_data - data)
import numpy as np
from scipy.stats import invgamma, norm
from ..field import Field
from ..sugar import makeOp
from ..multi import MultiField
from ..model import Model
from ..operators import SelectionOperator
from ..utilities import memo
class PointSources(Model):
def __init__(self, position, alpha, q):
super(PointSources, self).__init__(position)
self._alpha = alpha
self._q = q
def at(self, position):
return self.__class__(position, self._alpha, self._q)
@property
@memo
def value(self):
points = self.position['points'].to_global_data()
points = np.clip(points, None, 8.2)
points = Field(self.position['points'].domain, points)
return self.IG(points, self._alpha, self._q)
@property
@memo
def gradient(self):
u = self.position['points']
inner = norm.pdf(u.val)
outer_inv = invgamma.pdf(invgamma.ppf(norm.cdf(u.val),
self._alpha,
scale=self._q),
self._alpha, scale=self._q)
# FIXME
outer_inv = np.clip(outer_inv, 1e-20, None)
outer = 1/outer_inv
grad = Field(u.domain, val=inner*outer)
grad = makeOp(MultiField({'points': grad}))
return SelectionOperator(grad.target, 'points')*grad
@staticmethod
def IG(field, alpha, q):
foo = invgamma.ppf(norm.cdf(field.val), alpha, scale=q)
return Field(field.domain, val=foo)
@staticmethod
def IG_prime(field, alpha, q):
inner = norm.pdf(field.val)
outer = invgamma.pdf(invgamma.ppf(norm.cdf(field.val), alpha, scale=q), alpha, scale=q)
# # FIXME
# outer = np.clip(outer, 1e-20, None)
outer = 1/outer
return Field(field.domain, val=inner*outer)
@staticmethod
def inverseIG(u, alpha, q):
res = norm.ppf(invgamma.cdf(u, alpha, scale=q))
# # FIXME
# res = np.clip(res, 0, None)
return res
def make_smooth_sky_model(s_space, amplitude_model):
'''
Method for construction of correlated sky model
Parameters
----------
s_space : domain of sky model
amplitude_model : model for correlation structure
'''
from .. import (FFTOperator, Field, MultiField, PointwiseExponential,
PowerDistributor, Variable)
h_space = s_space.get_default_codomain()
ht = FFTOperator(h_space, s_space)
p_space = amplitude_model.value.domain[0]
power_distributor = PowerDistributor(h_space, p_space)
position = {}
position['xi'] = Field.from_random('normal', h_space)
position['tau'] = amplitude_model.position['tau']
position['phi'] = amplitude_model.position['phi']
position = MultiField(position)
xi = Variable(position)['xi']
A = power_distributor(amplitude_model)
logsky_h = A * xi
logsky = ht(logsky_h)
internals = {'logsky_h': logsky_h,
'power_distributor': power_distributor,
'ht': ht}
return PointwiseExponential(logsky), internals
def make_smooth_mf_sky_model(s_space_spatial, s_space_energy,
amplitude_model_spatial, amplitude_model_energy):
'''
Method for construction of correlated sky model
Parameters
----------
s_space : domain of sky model
amplitude_model : model for correlation structure
'''
from .. import (DomainTuple, Field, HarmonicTransformOperator, MultiField,
PointwiseExponential, PowerDistributor, Variable)
from ..linear_operators import DomainDistributor
h_space_spatial = s_space_spatial.get_default_codomain()
h_space_energy = s_space_energy.get_default_codomain()
h_space = DomainTuple.make((h_space_spatial, h_space_energy))
ht1 = HarmonicTransformOperator(h_space, space=0)
ht2 = HarmonicTransformOperator(ht1.target, space=1)
ht = ht2*ht1
p_space_spatial = amplitude_model_spatial.value.domain[0]
p_space_energy = amplitude_model_energy.value.domain[0]
pd_spatial = PowerDistributor(h_space, p_space_spatial, 0)
pd_energy = PowerDistributor(pd_spatial.domain, p_space_energy, 1)
pd = pd_spatial*pd_energy
dom_distr_0 = DomainDistributor(pd.domain, 0)
dom_distr_1 = DomainDistributor(pd.domain, 1)
a_spatial = dom_distr_1(amplitude_model_spatial)
a_energy = dom_distr_0(amplitude_model_energy)
a = a_spatial*a_energy
A = pd(a)
position = MultiField({'xi': Field.from_random('normal', h_space)})
xi = Variable(position)['xi']
logsky_h = A*xi
logsky = ht(logsky_h)
return PointwiseExponential(logsky)
from .diagonal_operator import DiagonalOperator
from .dof_distributor import DOFDistributor
from .endomorphic_operator import EndomorphicOperator
from .exp_transform import ExpTransform
from .fft_operator import FFTOperator
from .fft_smoothing_operator import FFTSmoothingOperator
from .geometry_remover import GeometryRemover
......@@ -10,15 +11,20 @@ from .laplace_operator import LaplaceOperator
from .linear_operator import LinearOperator
from .multi_adaptor import MultiAdaptor
from .power_distributor import PowerDistributor
from .qht_operator import QHTOperator
from .sampling_enabler import SamplingEnabler
from .sandwich_operator import SandwichOperator
from .scaling_operator import ScalingOperator
from .selection_operator import SelectionOperator
from .slope_operator import SlopeOperator
from .smoothness_operator import SmoothnessOperator
from .symmetrizing_operator import SymmetrizingOperator
__all__ = ["LinearOperator", "EndomorphicOperator", "ScalingOperator",
"DiagonalOperator", "HarmonicTransformOperator", "FFTOperator",
"FFTSmoothingOperator", "GeometryRemover",
"LaplaceOperator", "SmoothnessOperator", "PowerDistributor",
"InversionEnabler", "SandwichOperator", "SamplingEnabler",
"DOFDistributor", "SelectionOperator", "MultiAdaptor"]
"DOFDistributor", "SelectionOperator", "MultiAdaptor",
"ExpTransform", "SymmetrizingOperator", "QHTOperator",
"SlopeOperator"]
import numpy as np
from ..domain_tuple import DomainTuple
from ..domains import PowerSpace, RGSpace
from ..field import Field
from .linear_operator import LinearOperator
class ExpTransform(LinearOperator):
def __init__(self, target, dof):
if not ((isinstance(target, RGSpace) and target.harmonic) or
isinstance(target, PowerSpace)):
raise ValueError(
"Target must be a harmonic RGSpace or a power space.")
if np.isscalar(dof):
dof = np.full(len(target.shape), int(dof), dtype=np.int)
dof = np.array(dof)
ndim = len(target.shape)
t_mins = np.empty(ndim)
bindistances = np.empty(ndim)
self._bindex = [None] * ndim
self._frac = [None] * ndim
for i in range(ndim):
if isinstance(target, RGSpace):
rng = np.arange(target.shape[i])
tmp = np.minimum(rng, target.shape[i] + 1 - rng)
k_array = tmp * target.distances[i]
else:
k_array = target.k_lengths
# avoid taking log of first entry
log_k_array = np.log(k_array[1:])
# Interpolate log_k_array linearly
t_max = np.max(log_k_array)
t_min = np.min(log_k_array)
# Save t_min for later
t_mins[i] = t_min
bindistances[i] = (t_max - t_min) / (dof[i] - 1)
coord = np.append(0., 1. + (log_k_array - t_min) / bindistances[i])
self._bindex[i] = np.floor(coord).astype(int)
# Interpolated value is computed via
# (1.-frac)*<value from this bin> + frac*<value from next bin>
# 0 <= frac < 1.
self._frac[i] = coord - self._bindex[i]
from ..domains import LogRGSpace
log_space = LogRGSpace(2 * dof + 1, bindistances,
t_mins, harmonic=False)
self._target = DomainTuple.make(target)
self._domain = DomainTuple.make(log_space)
@property
def domain(self):
return self._domain
@property
def target(self):
return self._target
def apply(self, x, mode):
self._check_input(x, mode)
x = x.val
ndim = len(self.target.shape)
idx = ()
for d in range(ndim):
fst_dims = (1,) * d
lst_dims = (1,) * (ndim - d - 1)
wgt = self._frac[d].reshape(fst_dims + (-1,) + lst_dims)
# ADJOINT_TIMES
if mode == self.ADJOINT_TIMES:
shp = list(x.shape)
shp[d] = self._tgt(mode).shape[d]
xnew = np.zeros(shp, dtype=x.dtype)
np.add.at(xnew, idx + (self._bindex[d],), x * (1. - wgt))
np.add.at(xnew, idx + (self._bindex[d] + 1,), x * wgt)
# TIMES
else:
xnew = x[idx + (self._bindex[d],)] * (1. - wgt)
xnew += x[idx + (self._bindex[d] + 1,)] * wgt
x = xnew
idx = (slice(None),) + idx
return Field(self._tgt(mode), val=x)
@property
def capability(self):
return self.TIMES | self.ADJOINT_TIMES
from ..domain_tuple import DomainTuple
from ..field import Field
from ..utilities import hartley
from .linear_operator import LinearOperator
class QHTOperator(LinearOperator):
def __init__(self, domain, target):
if not domain.harmonic:
raise TypeError(
"HarmonicTransformOperator only works on a harmonic space")
if target.harmonic:
raise TypeError("Target is not a codomain of domain")
from ..domains import LogRGSpace
if not isinstance(domain, LogRGSpace):
raise ValueError("Domain has to be a LogRGSpace!")
if not isinstance(target, LogRGSpace):
raise ValueError("Target has to be a LogRGSpace!")
self._domain = DomainTuple.make(domain)
self._target = DomainTuple.make(target)
@property
def domain(self):
return self._domain
@property
def target(self):
return self._target
def apply(self, x, mode):
self._check_input(x, mode)
x = x.val * self.domain[0].scalar_dvol()
n = len(self.domain[0].shape)
rng = range(n) if mode == self.TIMES else reversed(range(n))
for i in rng:
sl = (slice(None),)*i + (slice(1, None),)
x[sl] = hartley(x[sl], axes=(i,))
return Field(self._tgt(mode), val=x)
@property
def capability(self):
return self.TIMES | self.ADJOINT_TIMES
import numpy as np
from ..domain_tuple import DomainTuple
from ..field import Field
from .linear_operator import LinearOperator
class SlopeOperator(LinearOperator):
def __init__(self, domain, target, sigmas):
self._domain = DomainTuple.make(domain)
self._target = DomainTuple.make(target)
if self.domain[0].shape != (len(self.target[0].shape) + 1,):
raise AssertionError("Shape mismatch!")
self._sigmas = sigmas
self.ndim = len(self.target[0].shape)
self.pos = np.zeros((self.ndim,) + self.target[0].shape)
if self.ndim == 1:
self.pos[0] = self.target[0].get_k_length_array().val
else:
shape = self.target[0].shape
for i in range(self.ndim):
rng = np.arange(target.shape[i])
tmp = np.minimum(
rng, target.shape[i] + 1 - rng) * target.bindistances[i]
fst_dims = (1,) * i
lst_dims = (1,) * (self.ndim - i - 1)
self.pos[i] += tmp.reshape(fst_dims + (shape[i],) + lst_dims)
@property
def sigmas(self):
return self._sigmas
@property
def domain(self):
return self._domain
@property
def target(self):
return self._target
def apply(self, x, mode):
self._check_input(x, mode)
# Times
if mode == self.TIMES:
inp = x.val
res = self.sigmas[-1] * inp[-1]
for i in range(self.ndim):
res += self.sigmas[i] * inp[i] * self.pos[i]
return Field(self.target, val=res)
# Adjoint times
res = np.zeros(self.domain[0].shape)
res[-1] = np.sum(x.val) * self.sigmas[-1]
for i in range(self.ndim):
res[i] = np.sum(self.pos[i] * x.val) * self.sigmas[i]
return Field(self.domain, val=res)
@property
def capability(self):