Skip to content
Snippets Groups Projects
Commit a7b8ea8b authored by Philipp Arras's avatar Philipp Arras
Browse files

Compile with Linearization (broken)

parent ce091e15
No related branches found
No related tags found
1 merge request!37Draft: Observation: introduce __hash__
Hide whitespace changes
Inline Side-by-side
resolvelib/devel.py
+ 25
− 5
View file @ a7b8ea8b
......@@ -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
Compare 5d3e4e57...deea702b
Subproject commit 5d3e4e577659ea5eb41a67e82f0cf4380ef599d1
Subproject commit deea702b488f7d610c60eadcbb83f087949a0496
resolvelib/resolvelib.cc
+ 89
− 46
View file @ a7b8ea8b
......@@ -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_;
};
return Linearization<py::dict,py::array>(apply(loc_), ftimes, fadjtimes);
}
};
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);
}
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 0 → 100644
+ 32
− 0
View file @ a7b8ea8b
// 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;
}
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment