diff --git a/nifty5/library/correlated_fields.py b/nifty5/library/correlated_fields.py index 04c2ce2863f8613fab256be6b5093204465be1fc..2f3272cc69f4e155717d753a45b97f5701df9fa6 100644 --- a/nifty5/library/correlated_fields.py +++ b/nifty5/library/correlated_fields.py @@ -185,9 +185,30 @@ class _SpecialSum(EndomorphicOperator): return self._contractor.adjoint(self._contractor(x)) +class _slice_extractor(LinearOperator): + #FIXME it should be tested if the the domain and target are consistent with the slice + def __init__(self, domain, target, sl): + self._domain = makeDomain(domain) + self._target = makeDomain(target) + self._sl = sl + self._capability = self.TIMES | self.ADJOINT_TIMES + + def apply(self, x, mode): + self._check_input(x, mode) + x = x.to_global_data() + if mode == self.TIMES: + res = x[self._sl] + res = res.reshape(self._target.shape) + else: + res = np.zeros(self._domain.shape) + res[self._sl] = x + return from_global_data(self._tgt(mode), res) + + class CorrelatedFieldMaker: def __init__(self): self._amplitudes = [] + self._spaces = [] def add_fluctuations_from_ops(self, target, fluctuations, flexibility, asperity, loglogavgslope, key, space = 0): @@ -260,6 +281,7 @@ class CorrelatedFieldMaker: # End move to tests self._amplitudes.append(ampl) + self._spaces.append(space) def add_fluctuations(self, target, fluctuations_mean, fluctuations_stddev, flexibility_mean, flexibility_stddev, asperity_mean, @@ -268,13 +290,13 @@ class CorrelatedFieldMaker: prefix = str(prefix) fluct = _lognormal_moment_matching(fluctuations_mean, fluctuations_stddev, - prefix + 'fluctuations', space) + prefix + 'fluctuations', space = space) flex = _lognormal_moment_matching(flexibility_mean, flexibility_stddev, - prefix + 'flexibility', space) + prefix + 'flexibility', space = space) asp = _lognormal_moment_matching(asperity_mean, asperity_stddev, - prefix + 'asperity', space) + prefix + 'asperity', space = space) avgsl = _normal(loglogavgslope_mean, loglogavgslope_stddev, - prefix + 'loglogavgslope', space) + prefix + 'loglogavgslope', space = space) self.add_fluctuations_from_ops(target, fluct, flex, asp, avgsl, prefix + 'spectrum', space) @@ -289,29 +311,40 @@ class CorrelatedFieldMaker: """ offset vs zeromode: volume factor """ + prefix = str(prefix) if offset is not None: offset = float(offset) - #TODO correct hspace - hspace = makeDomain( - [dd.target[0].harmonic_partner for dd in self._amplitudes]) + + hspace = [] + zeroind = () + for amp, space in zip(self._amplitudes, self._spaces): + dd = list(amp.target) + dd[space] = dd[space].harmonic_partner + hspace.extend(dd) + zeroind += (slice(None),)*space + (0,)*len(dd[space].shape) + hspace = makeDomain(hspace) + spaces = np.cumsum(self._spaces) + np.arange(len(self._spaces)) azm = _lognormal_moment_matching(offset_amplitude_mean, offset_amplitude_stddev, - prefix + 'zeromode') + prefix + 'zeromode', space = space) + foo = np.ones(hspace.shape) - zeroind = len(hspace.shape)*(0,) foo[zeroind] = 0 - azm = Adder(from_global_data(hspace, foo)) @ ValueInserter( - hspace, zeroind) @ azm - ht = HarmonicTransformOperator(hspace, space=0) - pd = PowerDistributor(hspace, self._amplitudes[0].target[0], 0) + ZeroModeInserter = _slice_extractor(hspace, azm.target, zeroind).adjoint + + azm = Adder(from_global_data(hspace, foo)) @ ZeroModeInserter @ azm + + #NOTE ht and pd operator able to act on several spaces might be nice + ht = HarmonicTransformOperator(hspace, space = spaces[0]) + pd = PowerDistributor(hspace, + self._amplitudes[0].target[spaces[0]], spaces[0]) for i in range(1, len(self._amplitudes)): - ht = HarmonicTransformOperator(ht.target, space=i) @ ht - pd = pd @ PowerDistributor( - pd.domain, self._amplitudes[i].target[0], space=i) + ht = HarmonicTransformOperator(ht.target, space = spaces[i]) @ ht + pd = pd @ PowerDistributor( pd.domain, + self._amplitudes[i].target[spaces[i]], space = spaces[i]) - spaces = tuple(range(len(self._amplitudes))) a = ContractionOperator(pd.domain, spaces[1:]).adjoint(self._amplitudes[0]) for i in range(1, len(self._amplitudes)):