start extending the Alm class to m subsets

python/alm.h

@@ -24,7 +24,6 @@
 #ifndef DUCC0_ALM_H
 #define DUCC0_ALM_H
 
-#if 1
 #include <complex>
 #include <cmath>
 #include "ducc0/infra/threading.h"
@@ -44,7 +43,9 @@ using namespace std;
 class Alm_Base
   {
   protected:
-    size_t lmax, mmax, tval;
+    size_t lmax;
+    vector<size_t> mval;
+    vector<ptrdiff_t> mstart;
 
   public:
     /*! Returns the total number of coefficients for maximum quantum numbers
@@ -56,27 +57,46 @@ class Alm_Base
       }
 
     /*! 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) {}
+    Alm_Base (size_t lmax_, size_t mmax_)
+      : lmax(lmax_), mval(mmax_+1), mstart(mmax_+1)
+      {
+      ptrdiff_t idx = 0;
+      for (size_t m=0; m<=mmax_; ++m)
+        {
+        mval[m] = m;
+        mstart[m] = idx-m;
+        idx += lmax-m+1;
+        }
+      }
+
+    Alm_Base (size_t lmax_, const vector<size_t> &mval_,
+              const vector<ptrdiff_t> &mstart_)
+      : lmax(lmax_), mval(mval_)
+      {
+      for (size_t i=0; i<mval.size(); ++i)
+        {
+        MR_assert(mval[i]<=lmax, "m >= lmax");
+        if (i>0)
+          MR_assert(mval[i]>mval[i-1], "m not strictly ascending")'
+        }
+      mstart.resize(mval.back(),-2*lmax);
+      for (size_t i=0; i<mval.size(); ++i)
+        mstart[mval[i]] = mstart_[i];
+      }
 
     /*! Returns the maximum \a l */
     size_t Lmax() const { return lmax; }
     /*! Returns the maximum \a m */
-    size_t Mmax() const { return mmax; }
+    size_t Mmax() const { return mval.back(); }
 
     /*! 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); }
+      { return mstart[m]; }
 
     /*! 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. */
@@ -87,21 +107,21 @@ template<typename T> class Alm: public Alm_Base
 
     template<typename Func> void applyLM(Func func)
       {
-      for (size_t m=0; m<=mmax; ++m)
+      for (auto m: mval)
         for (size_t l=m; l<=lmax; ++l)
-          func(l,m,alm.v(l,m));
+          func(l,m,alm.v(index(l,m)));
       }
 
   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"); }
+      { 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"); }
+      { 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)}) {}
+      : Alm_Base(lmax_,mmax_), alm ({Num_Alms(lmax,mmax_)}) {}
 
     /*! Sets all coefficients to zero. */
     void SetToZero ()
@@ -120,13 +140,20 @@ template<typename T> class Alm: public Alm_Base
     /*! \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),
+      MR_assert(factor.size()>size_t(Mmax()),
         "alm.ScaleM: factor array too short");
       applyLM([&factor](size_t /*l*/, size_t m, T &v){v*=factor(m);}); 
       }
     /*! Adds \a num to a_00. */
     template<typename T2> void Add (const T2 &num)
-      { alm.v(0)+=num; }
+      {
+      for (auto m: mval)
+        if (m==0)
+          {
+          alm.v(index_l0(m))+=num;
+          return;
+          }
+      }
 
     /*! Returns a reference to the specified coefficient. */
     T &operator() (size_t l, size_t m)
@@ -156,7 +183,7 @@ template<typename T> class Alm: public Alm_Base
     /*! Adds all coefficients from \a other to the own coefficients. */
     void Add (const Alm &other)
       {
-      MR_assert (conformable(other), "A_lm are not conformable");
+      MR_assert (Alm_Base(*this)==Alm_Base(other), "A_lm are not conformable");
       for (size_t m=0; m<alm.size(); ++m)
         alm.v(m) += other.alm(m);
       }

src/ducc0/not_yet_integrated/enumerate.h

+/*
+ *  This file is part of the MR utility library.
+ *
+ *  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
+ */
+
+/* Copyright (C) 2019-2020 Max-Planck-Society
+   Author: Nathan Reed */
+
+#ifndef DUCC0_ENUMERATE_H
+#define DUCC0_ENUMERATE_H
+
+#include <tuple>
+
+namespace mr {
+
+// copied from http://reedbeta.com/blog/python-like-enumerate-in-cpp17/
+template <typename T,
+          typename TIter = decltype(std::begin(std::declval<T>())),
+          typename = decltype(std::end(std::declval<T>()))>
+constexpr auto enumerate(T && iterable)
+  {
+  struct iterator
+    {
+    size_t i;
+    TIter iter;
+    bool operator != (const iterator & other) const { return iter != other.iter; }
+    void operator ++ () { ++i; ++iter; }
+    auto operator * () const { return std::tie(i, *iter); }
+    };
+  struct iterable_wrapper
+    {
+    T iterable;
+    auto begin() { return iterator{ 0, std::begin(iterable) }; }
+    auto end() { return iterator{ 0, std::end(iterable) }; }
+    };
+  return iterable_wrapper{ std::forward<T>(iterable) };
+  }
+
+}
+
+#endif