alm.h 8.98 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
/*
 *  This code is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This code is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this code; if not, write to the Free Software
 *  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 */

17
18
19
20
21
22
23
/*! \file alm.h
 *  Class for storing spherical harmonic coefficients.
 *
 *  Copyright (C) 2003-2020 Max-Planck-Society
 *  \author Martin Reinecke
 */

Martin Reinecke's avatar
Martin Reinecke committed
24
25
#ifndef DUCC0_ALM_H
#define DUCC0_ALM_H
26

27
#if 1
Martin Reinecke's avatar
Martin Reinecke committed
28
29
#include <complex>
#include <cmath>
Martin Reinecke's avatar
Martin Reinecke committed
30
#include "ducc0/infra/threading.h"
Martin Reinecke's avatar
Martin Reinecke committed
31
32
#endif

Martin Reinecke's avatar
Martin Reinecke committed
33
34
#include "ducc0/infra/mav.h"
#include "ducc0/infra/error_handling.h"
35

Martin Reinecke's avatar
Martin Reinecke committed
36
namespace ducc0 {
37

Martin Reinecke's avatar
Martin Reinecke committed
38
39
40
41
namespace detail_alm {

using namespace std;

42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
/*! Base class for calculating the storage layout of spherical harmonic
    coefficients. */
class Alm_Base
  {
  protected:
    size_t lmax, mmax, tval;

  public:
    /*! Returns the total number of coefficients for maximum quantum numbers
        \a l and \a m. */
    static size_t Num_Alms (size_t l, size_t m)
      {
      MR_assert(m<=l,"mmax must not be larger than lmax");
      return ((m+1)*(m+2))/2 + (m+1)*(l-m);
      }

    /*! Constructs an Alm_Base object with given \a lmax and \a mmax. */
    Alm_Base (size_t lmax_=0, size_t mmax_=0)
      : lmax(lmax_), mmax(mmax_), tval(2*lmax+1) {}

    /*! Returns the maximum \a l */
    size_t Lmax() const { return lmax; }
    /*! Returns the maximum \a m */
    size_t Mmax() const { return mmax; }

    /*! Returns an array index for a given m, from which the index of a_lm
        can be obtained by adding l. */
    size_t index_l0 (size_t m) const
      { return ((m*(tval-m))>>1); }

    /*! Returns the array index of the specified coefficient. */
    size_t index (size_t l, size_t m) const
      { return index_l0(m) + l; }

    /*! Returns \a true, if both objects have the same \a lmax and \a mmax,
        else  \a false. */
    bool conformable (const Alm_Base &other) const
      { return ((lmax==other.lmax) && (mmax==other.mmax)); }
  };

/*! Class for storing spherical harmonic coefficients. */
template<typename T> class Alm: public Alm_Base
  {
  private:
    mav<T,1> alm;

Martin Reinecke's avatar
Martin Reinecke committed
88
89
90
91
92
93
94
    template<typename Func> void applyLM(Func func)
      {
      for (size_t m=0; m<=mmax; ++m)
        for (size_t l=m; l<=lmax; ++l)
          func(l,m,alm.v(l,m));
      }

95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
  public:
    /*! Constructs an Alm object with given \a lmax and \a mmax. */
    Alm (mav<T,1> &data, size_t lmax_, size_t mmax_)
      : Alm_Base(lmax_, mmax_), alm(data)
      { MR_assert(alm.size()==Num_Alms(lmax, mmax), "bad array size"); }
    Alm (const mav<T,1> &data, size_t lmax_, size_t mmax_)
      : Alm_Base(lmax_, mmax_), alm(data)
      { MR_assert(alm.size()==Num_Alms(lmax, mmax), "bad array size"); }
    Alm (size_t lmax_=0, size_t mmax_=0)
      : Alm_Base(lmax_,mmax_), alm ({Num_Alms(lmax,mmax)}) {}

    /*! Sets all coefficients to zero. */
    void SetToZero ()
      { alm.fill(0); }

    /*! Multiplies all coefficients by \a factor. */
    template<typename T2> void Scale (const T2 &factor)
      { for (size_t m=0; m<alm.size(); ++m) alm.v(m)*=factor; }
    /*! \a a(l,m) *= \a factor[l] for all \a l,m. */
    template<typename T2> void ScaleL (const mav<T2,1> &factor)
      {
      MR_assert(factor.size()>size_t(lmax),
        "alm.ScaleL: factor array too short");
Martin Reinecke's avatar
Martin Reinecke committed
118
      applyLM([&factor](size_t l, size_t /*m*/, T &v){v*=factor(l);}); 
119
120
121
122
123
124
      }
    /*! \a a(l,m) *= \a factor[m] for all \a l,m. */
    template<typename T2> void ScaleM (const mav<T2,1> &factor)
      {
      MR_assert(factor.size()>size_t(mmax),
        "alm.ScaleM: factor array too short");
Martin Reinecke's avatar
Martin Reinecke committed
125
      applyLM([&factor](size_t /*l*/, size_t m, T &v){v*=factor(m);}); 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
      }
    /*! Adds \a num to a_00. */
    template<typename T2> void Add (const T2 &num)
      { alm.v(0)+=num; }

    /*! Returns a reference to the specified coefficient. */
    T &operator() (size_t l, size_t m)
      { return alm.v(index(l,m)); }
    /*! Returns a constant reference to the specified coefficient. */
    const T &operator() (size_t l, size_t m) const
      { return alm(index(l,m)); }

    /*! Returns a pointer for a given m, from which the address of a_lm
        can be obtained by adding l. */
    T *mstart (size_t m)
      { return &alm.v(index_l0(m)); }
    /*! Returns a pointer for a given m, from which the address of a_lm
        can be obtained by adding l. */
    const T *mstart (size_t m) const
      { return &alm(index_l0(m)); }

    /*! Returns a constant reference to the a_lm data. */
    const mav<T,1> &Alms() const { return alm; }

150
151
152
    /*! Returns a reference to the a_lm data. */
    mav<T,1> &Alms() { return alm; }

153
154
155
156
157
158
159
160
161
162
163
164
    ptrdiff_t stride() const
      { return alm.stride(0); }

    /*! Adds all coefficients from \a other to the own coefficients. */
    void Add (const Alm &other)
      {
      MR_assert (conformable(other), "A_lm are not conformable");
      for (size_t m=0; m<alm.size(); ++m)
        alm.v(m) += other.alm(m);
      }
  };

165
#if 1
Martin Reinecke's avatar
Martin Reinecke committed
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
/*! Class for calculation of the Wigner matrix at arbitrary angles, using Risbo
    recursion in a way that can be OpenMP-parallelised. This approach uses more
    memory and is slightly slower than wigner_d_risbo_scalar. */
class wigner_d_risbo_openmp
  {
  private:
    double p,q;
    vector<double> sqt;
    mav<double,2> d, dd;
    ptrdiff_t n;

  public:
    wigner_d_risbo_openmp(size_t lmax, double ang)
      : p(sin(ang/2)), q(cos(ang/2)), sqt(2*lmax+1),
        d({lmax+1,2*lmax+1}), dd({lmax+1,2*lmax+1}), n(-1)
181
      { for (size_t m=0; m<sqt.size(); ++m) sqt[m] = std::sqrt(double(m)); }
Martin Reinecke's avatar
Martin Reinecke committed
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231

    const mav<double,2> &recurse()
      {
      ++n;
      if (n==0)
        d.v(0,0) = 1;
      else if (n==1)
        {
        d.v(0,0) = q*q; d.v(0,1) = -p*q*sqt[2]; d.v(0,2) = p*p;
        d.v(1,0) = -d(0,1); d.v(1,1) = q*q-p*p; d.v(1,2) = d(0,1);
        }
      else
        {
        // padding
        int sign = (n&1)? -1 : 1;
        for (int i=0; i<=2*n-2; ++i)
          {
          d.v(n,i) = sign*d(n-2,2*n-2-i);
          sign=-sign;
          }
        for (int j=2*n-1; j<=2*n; ++j)
          {
          auto &xd((j&1) ? d : dd);
          auto &xdd((j&1) ? dd: d);
          double xj = 1./j;
          xdd.v(0,0) = q*xd(0,0);
          for (int i=1;i<j; ++i)
            xdd.v(0,i) = xj*sqt[j]*(q*sqt[j-i]*xd(0,i) - p*sqt[i]*xd(0,i-1));
          xdd.v(0,j) = -p*xd(0,j-1);
// parallelize
          for (int k=1; k<=n; ++k)
            {
            double t1 = xj*sqt[j-k]*q, t2 = xj*sqt[j-k]*p;
            double t3 = xj*sqt[k  ]*p, t4 = xj*sqt[k  ]*q;
            xdd.v(k,0) = xj*sqt[j]*(q*sqt[j-k]*xd(k,0) + p*sqt[k]*xd(k-1,0));
            for (int i=1; i<j; ++i)
              xdd.v(k,i) = t1*sqt[j-i]*xd(k,i) - t2*sqt[i]*xd(k,i-1)
                        + t3*sqt[j-i]*xd(k-1,i) + t4*sqt[i]*xd(k-1,i-1);
            xdd.v(k,j) = -t2*sqt[j]*xd(k,j-1) + t4*sqt[j]*xd(k-1,j-1);
            }
          }
        }
      return d;
      }
  };

template<typename T> void rotate_alm (Alm<complex<T>> &alm,
  double psi, double theta, double phi)
  {
  auto lmax=alm.Lmax();
Martin Reinecke's avatar
Martin Reinecke committed
232
233
  MR_assert (lmax==alm.Mmax(),
    "rotate_alm: lmax must be equal to mmax");
Martin Reinecke's avatar
Martin Reinecke committed
234

Martin Reinecke's avatar
Martin Reinecke committed
235
  if (theta!=0)
Martin Reinecke's avatar
Martin Reinecke committed
236
    {
Martin Reinecke's avatar
Martin Reinecke committed
237
238
239
240
241
242
243
244
245
246
    vector<complex<double> > exppsi(lmax+1), expphi(lmax+1);
    for (size_t m=0; m<=lmax; ++m)
      {
      exppsi[m] = polar(1.,-psi*m);
      expphi[m] = polar(1.,-phi*m);
      }
    vector<complex<double> > almtmp(lmax+1);
size_t nthreads=1;
    wigner_d_risbo_openmp rec(lmax,theta);
    for (size_t l=0; l<=lmax; ++l)
Martin Reinecke's avatar
Martin Reinecke committed
247
      {
Martin Reinecke's avatar
Martin Reinecke committed
248
      const auto &d(rec.recurse());
Martin Reinecke's avatar
Martin Reinecke committed
249

Martin Reinecke's avatar
Martin Reinecke committed
250
251
252
253
      for (size_t m=0; m<=l; ++m)
        almtmp[m] = complex<double>(alm(l,0))*d(l,l+m);

      execStatic(l+1, nthreads, 0, [&](Scheduler &sched)
Martin Reinecke's avatar
Martin Reinecke committed
254
        {
Martin Reinecke's avatar
Martin Reinecke committed
255
256
257
258
259
260
        auto rng=sched.getNext();
        auto lo=rng.lo;
        auto hi=rng.hi;

        bool flip = true;
        for (size_t mm=1; mm<=l; ++mm)
Martin Reinecke's avatar
Martin Reinecke committed
261
          {
Martin Reinecke's avatar
Martin Reinecke committed
262
263
264
265
266
267
268
269
270
271
272
          auto t1 = complex<double>(alm(l,mm))*exppsi[mm];
          bool flip2 = ((mm+lo)&1);
          for (auto m=lo; m<hi; ++m)
            {
            double d1 = flip2 ? -d(l-mm,l-m) : d(l-mm,l-m);
            double d2 = flip  ? -d(l-mm,l+m) : d(l-mm,l+m);
            double f1 = d1+d2, f2 = d1-d2;
            almtmp[m]+=complex<double>(t1.real()*f1,t1.imag()*f2);
            flip2 = !flip2;
            }
          flip = !flip;
Martin Reinecke's avatar
Martin Reinecke committed
273
          }
Martin Reinecke's avatar
Martin Reinecke committed
274
        });
Martin Reinecke's avatar
Martin Reinecke committed
275

Martin Reinecke's avatar
Martin Reinecke committed
276
277
278
279
280
281
282
283
284
285
286
287
      for (size_t m=0; m<=l; ++m)
        alm(l,m) = complex<T>(almtmp[m]*expphi[m]);
      }
    }
  else
    {
    for (size_t m=0; m<=lmax; ++m)
      {
      auto ang = polar(1.,-(psi+phi)*m);
      for (size_t l=m; l<=lmax; ++l)
        alm(l,m) *= ang;
      }
Martin Reinecke's avatar
Martin Reinecke committed
288
289
290
291
292
293
294
    }
  }
#endif
}

using detail_alm::Alm_Base;
using detail_alm::Alm;
295
#if 1
Martin Reinecke's avatar
Martin Reinecke committed
296
297
using detail_alm::rotate_alm;
#endif
298
299
300
}

#endif