Commit 47cbdda4 authored by Martin Reinecke's avatar Martin Reinecke
Browse files

Merge branch 'nifty6_einsum' into 'NIFTy_6'

Introduce a numpy.einsum wrapper with derivatives

See merge request !460
parents 38bac55e 9b11f80b
Pipeline #75095 passed with stages
in 36 minutes and 28 seconds
......@@ -25,6 +25,7 @@ from .operators.adder import Adder
from .operators.diagonal_operator import DiagonalOperator
from .operators.distributors import DOFDistributor, PowerDistributor
from .operators.domain_tuple_field_inserter import DomainTupleFieldInserter
from .operators.einsum import LinearEinsum, MultiLinearEinsum
from .operators.contraction_operator import ContractionOperator
from .operators.linear_interpolation import LinearInterpolator
from .operators.endomorphic_operator import EndomorphicOperator
......
# This program 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 3 of the License, or
# (at your option) any later version.
#
# This program 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 program. If not, see <http://www.gnu.org/licenses/>.
#
# Copyright(C) 2013-2019 Max-Planck-Society
# Authors: Gordian Edenhofer, Philipp Frank
#
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
import numpy as np
import string
from ..domain_tuple import DomainTuple
from ..linearization import Linearization
from ..field import Field
from ..multi_domain import MultiDomain
from ..multi_field import MultiField
from .operator import Operator
from .linear_operator import LinearOperator
class MultiLinearEinsum(Operator):
"""Multi-linear Einsum operator with corresponding derivates
Parameters
----------
domain : MultiDomain or dict{name: DomainTuple}
The operator's input domain.
subscripts : str
The subscripts which is passed to einsum.
key_order: tuple of str, optional
The order of the keys in the multi-field. If not specified, defaults to
the order of the keys in the multi-field.
static_mf: MultiField or dict{name: Field}, optional
A dictionary like type from which Fields are to be taken if the key from
`key_order` is not part of the `domain`. Fields in this object are
supposed to be static as they will not appear as FieldAdapter in the
Linearization.
optimize: bool, String or List, optional
Parameter passed on to einsum_path.
Notes
-----
By convention :class:`MultiLinearEinsum` only performs operations with
lower indices. Therefore no complex conjugation is performed on complex
inputs. To achieve operations with upper/lower indices use
:class:`PartialConjugate` before applying this operator.
"""
def __init__(self, domain, subscripts,
key_order=None, static_mf=None, optimize='optimal'):
self._domain = MultiDomain.make(domain)
if key_order is None:
self._key_order = tuple(self._domain.keys())
else:
self._key_order = key_order
if static_mf is not None and key_order is None:
ve = "`key_order` mus be specified if additional fields are munged"
raise ValueError(ve)
self._stat_mf = static_mf
iss, oss, *rest = subscripts.split("->")
iss_spl = iss.split(",")
len_consist = len(self._key_order) == len(iss_spl)
sscr_consist = all(o in iss for o in oss)
if rest or not sscr_consist or "," in oss or not len_consist:
raise ValueError(f"invalid subscripts specified; got {subscripts}")
ve = f"invalid order of keys {self._key_order} for subscripts {subscripts}"
shapes, numpy_subscripts, subscriptmap = {}, '', {}
alphabet = list(string.ascii_lowercase)[::-1]
for k, ss in zip(self._key_order, iss_spl):
dom = self._domain[k] if k in self._domain.keys(
) else self._stat_mf[k].domain
if len(dom) != len(ss):
raise ValueError(ve)
for i, a in enumerate(list(ss)):
if a not in subscriptmap.keys():
subscriptmap[a] = [alphabet.pop() for _ in
range(len(dom[i].shape))]
numpy_subscripts += ''.join(subscriptmap[a])
numpy_subscripts += ','
shapes[k] = dom.shape
numpy_subscripts = numpy_subscripts[:-1] + '->'
dom_sscr = dict(zip(self._key_order, iss_spl))
tgt = []
for o in oss:
k_hit = tuple(k for k, sscr in dom_sscr.items() if o in sscr)[0]
dom_k_idx = dom_sscr[k_hit].index(o)
if k_hit in self._domain.keys():
tgt += [self._domain[k_hit][dom_k_idx]]
else:
if k_hit not in self._stat_mf.keys():
ve = f"{k_hit} is not in domain nor in static_mf"
raise ValueError(ve)
tgt += [self._stat_mf[k_hit].domain[dom_k_idx]]
numpy_subscripts += ''.join(subscriptmap[o])
self._target = DomainTuple.make(tgt)
numpy_iss, numpy_oss, *_ = numpy_subscripts.split("->")
numpy_iss_spl = numpy_iss.split(",")
self._sscr_endswith = dict()
self._linpaths = dict()
for k, (i, ss) in zip(self._key_order, enumerate(numpy_iss_spl)):
left_ss_spl = (*numpy_iss_spl[:i], *numpy_iss_spl[i + 1:], ss)
linpath = '->'.join((','.join(left_ss_spl), numpy_oss))
plc = tuple(np.broadcast_to(np.nan, shapes[q]) for q in shapes if q != k)
plc += (np.broadcast_to(np.nan, shapes[k]),)
self._sscr_endswith[k] = linpath
self._linpaths[k] = np.einsum_path(linpath, *plc, optimize=optimize)[0]
if isinstance(optimize, list):
path = optimize
else:
plc = (np.broadcast_to(np.nan, shapes[k]) for k in shapes)
path = np.einsum_path(numpy_subscripts, *plc, optimize=optimize)[0]
self._sscr = numpy_subscripts
self._ein_kw = {"optimize": path}
def apply(self, x):
self._check_input(x)
if isinstance(x, Linearization):
val = x.val.val
else:
val = x.val
v = (
val[k] if k in val else self._stat_mf[k].val
for k in self._key_order
)
res = np.einsum(self._sscr, *v, **self._ein_kw)
if isinstance(x, Linearization):
jac = None
for wrt in self.domain.keys():
plc = {
k: x.val[k] if k in x.val else self._stat_mf[k]
for k in self._key_order if k != wrt
}
mf_wo_k = MultiField.from_dict(plc)
ss = self._sscr_endswith[wrt]
# Use the fact that the insertion order in a dictionary is the
# ordering of keys as to pass on `key_order`
jac_k = LinearEinsum(
self.domain[wrt],
mf_wo_k,
ss,
key_order=tuple(plc.keys()),
optimize=self._linpaths[wrt],
_target=self._target,
_calling_as_lin=True
).ducktape(wrt)
jac = jac + jac_k if jac is not None else jac_k
return x.new(Field.from_raw(self.target, res), jac)
return Field.from_raw(self.target, res)
class LinearEinsum(LinearOperator):
"""Linear Einsum operator with exactly one freely varying field
Parameters
----------
domain : Domain, DomainTuple or tuple of Domain
The operator's input domain.
mf : MultiField
The first part of the left-hand side of the einsum.
subscripts : str
The subscripts which is passed to einsum. Everything before the very
last scripts before the '->' is treated as part of the fixed mulfi-
field while the last scripts are taken to correspond to the freely
varying field.
key_order: tuple of str, optional
The order of the keys in the multi-field. If not specified, defaults to
the order of the keys in the multi-field.
optimize: bool, String or List, optional
Parameter passed on to einsum_path.
Notes
-----
By convention :class:`LinearEinsum` only performs operations with
lower indices. Therefore no complex conjugation is performed on complex
inputs or mf. To achieve operations with upper/lower indices use
:class:`PartialConjugate` before applying this operator.
"""
def __init__(self, domain, mf, subscripts, key_order=None, optimize='optimal',
_target=None, _calling_as_lin=False):
self._domain = DomainTuple.make(domain)
if _calling_as_lin:
self._init_wo_preproc(mf, subscripts, key_order, optimize, _target)
else:
self._mf = mf
if key_order is None:
_key_order = tuple(self._mf.domain.keys())
else:
_key_order = key_order
self._ein_kw = {"optimize": optimize}
iss, oss, *rest = subscripts.split("->")
iss_spl = iss.split(",")
sscr_consist = all(o in iss for o in oss)
len_consist = len(_key_order) == len(iss_spl[:-1])
if rest or not sscr_consist or "," in oss or not len_consist:
raise ValueError(f"invalid subscripts specified; got {subscripts}")
ve = f"invalid order of keys {_key_order} for subscripts {subscripts}"
shapes, numpy_subscripts, subscriptmap = (), '', {}
alphabet = list(string.ascii_lowercase)
for k, ss in zip(_key_order, iss_spl[:-1]):
dom = self._mf[k].domain
if len(dom) != len(ss):
raise ValueError(ve)
for i, a in enumerate(list(ss)):
if a not in subscriptmap.keys():
subscriptmap[a] = [alphabet.pop() for _ in
range(len(dom[i].shape))]
numpy_subscripts += ''.join(subscriptmap[a])
numpy_subscripts += ','
shapes += (dom.shape,)
if len(self._domain) != len(iss_spl[-1]):
raise ValueError(ve)
for i, a in enumerate(list(iss_spl[-1])):
if a not in subscriptmap.keys():
subscriptmap[a] = [alphabet.pop() for _ in
range(len(self._domain[i].shape))]
numpy_subscripts += ''.join(subscriptmap[a])
shapes += (self._domain.shape,)
numpy_subscripts += '->'
dom_sscr = dict(zip(_key_order, iss_spl[:-1]))
dom_sscr[id(self)] = iss_spl[-1]
tgt = []
for o in oss:
k_hit = tuple(k for k, sscr in dom_sscr.items() if o in sscr)[0]
dom_k_idx = dom_sscr[k_hit].index(o)
if k_hit in _key_order:
tgt += [self._mf.domain[k_hit][dom_k_idx]]
else:
assert k_hit == id(self)
tgt += [self._domain[dom_k_idx]]
numpy_subscripts += "".join(subscriptmap[o])
_target = DomainTuple.make(tgt)
self._sscr = numpy_subscripts
if isinstance(optimize, list):
path = optimize
else:
plc = (np.broadcast_to(np.nan, shp) for shp in shapes)
path = np.einsum_path(numpy_subscripts, *plc, optimize=optimize)[0]
self._init_wo_preproc(mf, numpy_subscripts, _key_order, path, _target)
def _init_wo_preproc(self, mf, subscripts, keyorder, optimize, target):
self._ein_kw = {"optimize": optimize}
self._mf = mf
self._sscr = subscripts
self._key_order = keyorder
self._target = target
iss, oss, *_ = subscripts.split("->")
iss_spl = iss.split(",")
adj_iss = ",".join((",".join(iss_spl[:-1]), oss))
self._adj_sscr = "->".join((adj_iss, iss_spl[-1]))
self._capability = self.TIMES | self.ADJOINT_TIMES
def apply(self, x, mode):
self._check_input(x, mode)
if mode == self.TIMES:
dom, ss, mf = self.target, self._sscr, self._mf
else:
dom, ss, mf = self.domain, self._adj_sscr, self._mf.conjugate()
res = np.einsum(
ss, *(mf[k].val for k in self._key_order), x.val,
**self._ein_kw
)
return Field.from_raw(dom, res)
# This program 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 3 of the License, or
# (at your option) any later version.
#
# This program 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 program. If not, see <http://www.gnu.org/licenses/>.
#
# Copyright(C) 2013-2019 Max-Planck-Society
# Authors: Philipp Frank
#
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
from .endomorphic_operator import EndomorphicOperator
from ..multi_domain import MultiDomain
from ..multi_field import MultiField
class PartialConjugate(EndomorphicOperator):
"""Perform partial conjugation of a :class:`MultiField`
Parameters
----------
domain : MultiDomain
The operator's input domain and output target
conjugation_keys : iterable of string
The keys of the :class:`MultiField` for which complex conjugation
should be performed.
"""
def __init__(self, domain, conjugation_keys):
if not isinstance(domain, MultiDomain):
raise ValueError("MultiDomain expected!")
indom = (key in domain.keys() for key in conjugation_keys)
if sum(indom) != len(conjugation_keys):
raise ValueError("conjugation_keys not in domain!")
self._domain = domain
self._conjugation_keys = conjugation_keys
self._capabilities = self._all_ops
def apply(self, x, mode):
self._check_input(x, mode)
x = x.to_dict()
for k in self._conjugation_keys:
x[k] = x[k].conjugate()
return MultiField.from_dict(x, self._domain)
# This program 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 3 of the License, or
# (at your option) any later version.
#
# This program 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 program. If not, see <http://www.gnu.org/licenses/>.
#
# Copyright(C) 2013-2019 Max-Planck-Society
# Authors: Gordian Edenhofer, Philipp Frank
#
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
import pytest
from numpy.testing import assert_allclose
import numpy as np
from nifty6.extra import check_jacobian_consistency, consistency_check
import nifty6 as ift
from ..common import list2fixture, setup_function, teardown_function
spaces = (ift.UnstructuredDomain(4),
ift.RGSpace((3,2)),
ift.LMSpace(5),
ift.GLSpace(4))
space1 = list2fixture(spaces)
space2 = list2fixture(spaces)
dtype = list2fixture([np.float64, np.complex128])
def test_linear_einsum_outer(space1, space2, dtype, n_invocations=10):
mf_dom = ift.MultiDomain.make(
{
"dom01": space1,
"dom02":
ift.DomainTuple.make(
(space1, space2)
)
}
)
mf = ift.from_random("normal", mf_dom, dtype=dtype)
ss = "i,ij,j->ij"
key_order = ("dom01", "dom02")
le = ift.LinearEinsum(space2, mf, ss, key_order=key_order)
assert consistency_check(le, domain_dtype=dtype,target_dtype=dtype) is None
le_ift = ift.DiagonalOperator(
mf["dom01"], domain=mf_dom["dom02"], spaces=0
) @ ift.DiagonalOperator(mf["dom02"]) @ ift.OuterProduct(
ift.full(mf_dom["dom01"], 1.), ift.DomainTuple.make(mf_dom["dom02"][1])
)
for _ in range(n_invocations):
r = ift.from_random("normal", le.domain, dtype=dtype)
assert_allclose(le(r).val, le_ift(r).val)
r_adj = ift.from_random("normal", le.target, dtype=dtype)
assert_allclose(le.adjoint(r_adj).val, le_ift.adjoint(r_adj).val)
def test_linear_einsum_contraction(space1, space2, dtype, n_invocations=10):
mf_dom = ift.MultiDomain.make(
{
"dom01": space1,
"dom02":
ift.DomainTuple.make(
(space1, space2)
)
}
)
mf = ift.from_random("normal", mf_dom, dtype=dtype)
ss = "i,ij,j->i"
key_order = ("dom01", "dom02")
le = ift.LinearEinsum(space2, mf, ss, key_order=key_order)
assert consistency_check(le, domain_dtype=dtype,target_dtype=dtype) is None
le_ift = ift.ContractionOperator(mf_dom["dom02"], 1) @ ift.DiagonalOperator(
mf["dom01"], domain=mf_dom["dom02"], spaces=0
) @ ift.DiagonalOperator(mf["dom02"]) @ ift.OuterProduct(
ift.full(mf_dom["dom01"], 1.), ift.DomainTuple.make(mf_dom["dom02"][1])
)
for _ in range(n_invocations):
r = ift.from_random("normal", le.domain, dtype=dtype)
assert_allclose(le(r).val, le_ift(r).val)
r_adj = ift.from_random("normal", le.target, dtype=dtype)
assert_allclose(le.adjoint(r_adj).val, le_ift.adjoint(r_adj).val)
def test_multi_linear_einsum_outer(
space1, space2, dtype, n_invocations=10, ntries=100
):
mf_dom = ift.MultiDomain.make(
{
"dom01": space1,
"dom02":ift.DomainTuple.make((space1, space2)),
"dom03": space2
}
)
ss = "i,ij,j->ij"
key_order = ("dom01", "dom02", "dom03")
mle = ift.MultiLinearEinsum(mf_dom, ss, key_order=key_order)
check_jacobian_consistency(
mle, ift.from_random("normal", mle.domain, dtype=dtype), ntries=ntries
)
outer_i = ift.OuterProduct(
ift.full(mf_dom["dom03"], 1.), ift.DomainTuple.make(mf_dom["dom02"][0])
)
outer_j = ift.OuterProduct(
ift.full(mf_dom["dom01"], 1.), ift.DomainTuple.make(mf_dom["dom02"][1])
)
# SwitchSpacesOperator is equivalent to LinearEinsum with "ij->ji"
mle_ift = ift.SwitchSpacesOperator(
outer_i.target, 1
) @ outer_i @ ift.FieldAdapter(mf_dom["dom01"], "dom01") * ift.FieldAdapter(
mf_dom["dom02"], "dom02"
) * (outer_j @ ift.FieldAdapter(mf_dom["dom03"], "dom03"))
for _ in range(n_invocations):
rl = ift.Linearization.make_var(
ift.from_random("normal", mle.domain, dtype=dtype))
mle_rl, mle_ift_rl = mle(rl), mle_ift(rl)
assert_allclose(mle_rl.val.val, mle_ift_rl.val.val)
assert_allclose(mle_rl.jac(rl.val).val, mle_ift_rl.jac(rl.val).val)
rj_adj = ift.from_random("normal", mle_rl.jac.target, dtype=dtype)
mle_j_val = mle_rl.jac.adjoint(rj_adj).val
mle_ift_j_val = mle_ift_rl.jac.adjoint(rj_adj).val
for k in mle_ift.domain.keys():
assert_allclose(mle_j_val[k], mle_ift_j_val[k])
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment