Compile with Linearization (broken)

resolvelib/devel.py

@@ -23,8 +23,28 @@ from cpp2py import Pybind11Operator
 from time import time
 
 
+def operator_equality(op0, op1, ntries=20):
+    dom = op0.domain
+    assert op0.domain == op1.domain
+    assert op0.target == op1.target
+    for ii in range(ntries):
+        loc = ift.from_random(dom)
+        res0 = op0(loc)
+        res1 = op1(loc)
+        ift.extra.assert_allclose(res0, res1)
+
+        linloc = ift.Linearization.make_var(loc)
+        res0 = op0(linloc).jac(0.23*loc)
+        res1 = op1(linloc).jac(0.23*loc)
+        ift.extra.assert_allclose(res0, res1)
+
+    ift.extra.check_operator(op0, loc, ntries=ntries)
+    ift.extra.check_operator(op1, loc, ntries=ntries)
+
+
+
 pdom = rve.PolarizationSpace(["I", "Q", "U", "V"])
-sdom = ift.RGSpace([4000, 4000])
+sdom = ift.RGSpace([10, 10])
 
 dom = rve.default_sky_domain(pdom=pdom, sdom=sdom)
 dom = {kk: dom[1:] for kk in pdom.labels}
@@ -32,11 +52,11 @@ dom = {kk: dom[1:] for kk in pdom.labels}
 tgt = rve.default_sky_domain(pdom=pdom, sdom=sdom)
 
 opold = rve.polarization_matrix_exponential(tgt) @ rve.MultiFieldStacker(tgt, 0, tgt[0].labels)
+op = Pybind11Operator(dom, tgt, resolvelib.PolarizationMatrixExponential(1))
+
+operator_equality(opold, op)
+exit()
 
-nthreads = 1
-loc = ift.full(dom, 1.2)
-t0 = time()
-res0 = opold(loc)
 print(f"Old implementation: {(time()-t0):.2f} s")
 ntries = 50
 for nthreads  in [8]:

ducc @ deea702b

-Subproject commit 5d3e4e577659ea5eb41a67e82f0cf4380ef599d1
+Subproject commit deea702b488f7d610c60eadcbb83f087949a0496

resolvelib/resolvelib.cc

@@ -24,14 +24,14 @@
 #include <pybind11/pybind11.h>
 
 #include "ducc0/bindings/pybind_utils.h"
-#include "ducc0/infra/threading.h"
 using namespace pybind11::literals;
 namespace py = pybind11;
 auto None = py::none();
 // Includes related to pybind11
 
-using namespace std;
+#include "shape_helper.h"
 
+using namespace std;
 
 template<typename Tin, typename Tout>
 class Linearization {
@@ -51,6 +51,18 @@ class Linearization {
     const function<Tin (const Tout &)> jat;
 };
 
+template<typename Tin, typename Tout>
+void add_linearization(py::module_ &msup, const char *name) {
+    py::class_<Linearization<Tin, Tout>>(msup, name)
+       .def(py::init<const Tout &,
+                     function<Tout(const Tin &)>,
+                     function<Tin (const Tout &)>
+                     >())
+       .def("position",          &Linearization<Tin, Tout>::position)
+       .def("jac_times",         &Linearization<Tin, Tout>::jac_times)
+       .def("jac_adjoint_times", &Linearization<Tin, Tout>::jac_adjoint_times);
+}
+
 
 template<typename T, size_t ndim>
 class PolarizationMatrixExponential {
@@ -58,14 +70,16 @@ class PolarizationMatrixExponential {
     const size_t nthreads;
   public:
     PolarizationMatrixExponential(size_t nthreads_=1) : nthreads(nthreads_) {}
-    py::array apply(const py::dict &inp) const {
-      auto I {ducc0::to_cmav<T, ndim>(inp["I"])},
-           Q {ducc0::to_cmav<T, ndim>(inp["Q"])},
-           U {ducc0::to_cmav<T, ndim>(inp["U"])},
-           V {ducc0::to_cmav<T, ndim>(inp["V"])};
+    py::array apply(const py::dict &inp_) const {
+      // Parse input
+      auto I {ducc0::to_cmav<T, ndim>(inp_["I"])},
+           Q {ducc0::to_cmav<T, ndim>(inp_["Q"])},
+           U {ducc0::to_cmav<T, ndim>(inp_["U"])},
+           V {ducc0::to_cmav<T, ndim>(inp_["V"])};
+      // /Parse input
 
       // Instantiate output array
-      auto out_ = ducc0::make_Pyarr<T>({4, I.shape()[0], I.shape()[1], I.shape()[2], I.shape()[3]});
+      auto out_ = ducc0::make_Pyarr<T>(combine_shapes(4, I.shape()));
       auto out  = ducc0::to_vmav<T, ndim+1>(out_);
       auto outI = ducc0::subarray<ndim>(out, {{0}, {}, {}, {}, {}}),
            outQ = ducc0::subarray<ndim>(out, {{1}, {}, {}, {}, {}}),
@@ -73,14 +87,10 @@ class PolarizationMatrixExponential {
            outV = ducc0::subarray<ndim>(out, {{3}, {}, {}, {}, {}});
       // /Instantiate output array
 
-      ducc0::mav_apply([](const auto &ii,
-                          const auto &qq,
-                          const auto &uu,
-                          const auto &vv,
-                          auto &oii,
-                          auto &oqq,
-                          auto &ouu,
-                          auto &ovv
+      ducc0::mav_apply([](const auto &ii, const auto &qq,
+                          const auto &uu, const auto &vv,
+                          auto &oii, auto &oqq,
+                          auto &ouu, auto &ovv
                           ){
               auto pol0{qq*qq + uu*uu + vv*vv};
               auto pol1{-0.5*log(pol0)};
@@ -94,42 +104,75 @@ class PolarizationMatrixExponential {
       return out_;
     }
 
-    // Linearization<py::array,py::array> apply_with_jac(const py::dict &inp) {
-    //     function<py::array(const py::dict &)> f =
-    //       [=](const py::dict &inp2) {
-    //         MR_fail("Not implemented yet");
-    //     };
-    //     function<py::array(const py::dict &)> ftimes =
-    //         [=](const py::dict &inp2) { return f(inp2); };
-    //     function<py::dict(const py::array &)> fadjtimes =
-    //         [=](const py::array &inp2) { return f(inp2); };
-
-    //     return Linearization<py::dict,py::array> {apply(inp), ftimes, fadjtimes};
-    // }
+    Linearization<py::dict,py::array> apply_with_jac(const py::dict &loc_) {
+
+        function<py::array(const py::dict &)> ftimes =
+            [&](const py::dict &inp_) {
+              // Parse input
+              auto I {ducc0::to_cmav<T, ndim>(inp_["I"])},
+                   Q {ducc0::to_cmav<T, ndim>(inp_["Q"])},
+                   U {ducc0::to_cmav<T, ndim>(inp_["U"])},
+                   V {ducc0::to_cmav<T, ndim>(inp_["V"])};
+              // /Parse input
+
+              // Instantiate output array
+              auto out_ = ducc0::make_Pyarr<T>(combine_shapes(4, I.shape()));
+              auto out  = ducc0::to_vmav<T, ndim+1>(out_);
+              auto outI = ducc0::subarray<ndim>(out, {{0}, {}, {}, {}, {}}),
+                   outQ = ducc0::subarray<ndim>(out, {{1}, {}, {}, {}, {}}),
+                   outU = ducc0::subarray<ndim>(out, {{2}, {}, {}, {}, {}}),
+                   outV = ducc0::subarray<ndim>(out, {{3}, {}, {}, {}, {}});
+              // /Instantiate output array
+
+              return out_;
             };
 
+        function<py::dict(const py::array &)> fadjtimes =
+            [&](const py::array &inp_) {
+              // Parse input
+              auto inp {ducc0::to_cmav<T, ndim+1>(inp_)};
+              auto I {ducc0::subarray<ndim>(inp, {{0}, {}, {}, {}, {}})},
+                   Q {ducc0::subarray<ndim>(inp, {{1}, {}, {}, {}, {}})},
+                   U {ducc0::subarray<ndim>(inp, {{2}, {}, {}, {}, {}})},
+                   V {ducc0::subarray<ndim>(inp, {{3}, {}, {}, {}, {}})};
+              // /Parse input
+
+              // Instantiate output
+              auto outI_ = ducc0::make_Pyarr<T>(I.shape());
+              auto outQ_ = ducc0::make_Pyarr<T>(I.shape());
+              auto outU_ = ducc0::make_Pyarr<T>(I.shape());
+              auto outV_ = ducc0::make_Pyarr<T>(I.shape());
+
+              auto outI {ducc0::to_vmav<T, ndim>(outI_)},
+                   outQ {ducc0::to_vmav<T, ndim>(outQ_)},
+                   outU {ducc0::to_vmav<T, ndim>(outU_)},
+                   outV {ducc0::to_vmav<T, ndim>(outV_)};
+              // /Instantiate output
+
+              // Pack into dictionary
+              py::dict out_;
+              out_["I"] = outI_;
+              out_["Q"] = outQ_;
+              out_["U"] = outU_;
+              out_["V"] = outV_;
+              // /Pack into dictionary
+              return out_;
+            };
 
-template<typename Tin, typename Tout>
-void add_linearization(py::module_ &msup, const char *name) {
-    py::class_<Linearization<Tin, Tout>>(msup, name)
-       .def(py::init<const Tout &,
-                     function<Tout(const Tin &)>,
-                     function<Tin (const Tout &)>
-                     >())
-       .def("position",          &Linearization<Tin, Tout>::position)
-       .def("jac_times",         &Linearization<Tin, Tout>::jac_times)
-       .def("jac_adjoint_times", &Linearization<Tin, Tout>::jac_adjoint_times);
+        return Linearization<py::dict,py::array>(apply(loc_), ftimes, fadjtimes);
     }
+};
+
 
 
 PYBIND11_MODULE(resolvelib, m) {
     py::class_<PolarizationMatrixExponential<double, 4>>(m, "PolarizationMatrixExponential")
         .def(py::init<size_t>())
-        .def("apply", &PolarizationMatrixExponential<double, 4>::apply);
-        //.def("apply_with_jac", &PolarizationMatrixExponential::apply_with_jac);
+        .def("apply",          &PolarizationMatrixExponential<double, 4>::apply)
+        .def("apply_with_jac", &PolarizationMatrixExponential<double, 4>::apply_with_jac);
 
-    // add_linearization<py::array, py::array>(m, "Linearization_field2field");
-    // add_linearization<py::array, py::dict >(m, "Linearization_field2mfield");
-    // add_linearization<py::dict , py::array>(m, "Linearization_mfield2field");
-    // add_linearization<py::dict , py::dict >(m, "Linearization_mfield2mfield");
+    add_linearization<py::array, py::array>(m, "Linearization_field2field");
+    add_linearization<py::array, py::dict >(m, "Linearization_field2mfield");
+    add_linearization<py::dict , py::array>(m, "Linearization_mfield2field");
+    add_linearization<py::dict , py::dict >(m, "Linearization_mfield2mfield");
 }

resolvelib/shape_helper.h

+// Author: Martin Reinecke
+
+#include <array>
+#include <algorithm>
+
+using namespace std;
+
+
+template<typename T, size_t L1, size_t L2>
+  array<T,L1+L2> combine_shapes(const array<T, L1> &a1, const array<T, L2> &a2)
+  {
+  array<T, L1+L2> res;
+  copy_n(a1.begin(), L1, res.begin());
+  copy_n(a2.begin(), L2, res.begin()+L1);
+  return res;
+  }
+template<typename T, size_t L>
+  array<T,1+L> combine_shapes(size_t s1, const array<T, L> &a)
+  {
+  array<T, 1+L> res;
+  res[0] = s1;
+  copy_n(a.begin(), L, res.begin()+1);
+  return res;
+  }
+template<typename T, size_t L>
+  array<T,1+L> combine_shapes(const array<T, L> &a, size_t s2)
+  {
+  array<T, 1+L> res;
+  copy_n(a.begin(), L, res.begin());
+  res[L] = s2;
+  return res;
+  }