Kokkos Core Kernels Package Version of the Day
Loading...
Searching...
No Matches
Kokkos_Complex.hpp
1//@HEADER
2// ************************************************************************
3//
4// Kokkos v. 4.0
5// Copyright (2022) National Technology & Engineering
6// Solutions of Sandia, LLC (NTESS).
7//
8// Under the terms of Contract DE-NA0003525 with NTESS,
9// the U.S. Government retains certain rights in this software.
10//
11// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
12// See https://kokkos.org/LICENSE for license information.
13// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
14//
15//@HEADER
16#ifndef KOKKOS_COMPLEX_HPP
17#define KOKKOS_COMPLEX_HPP
18#ifndef KOKKOS_IMPL_PUBLIC_INCLUDE
19#define KOKKOS_IMPL_PUBLIC_INCLUDE
20#define KOKKOS_IMPL_PUBLIC_INCLUDE_NOTDEFINED_COMPLEX
21#endif
22
23#include <Kokkos_Atomic.hpp>
24#include <Kokkos_MathematicalFunctions.hpp>
25#include <Kokkos_NumericTraits.hpp>
26#include <Kokkos_ReductionIdentity.hpp>
27#include <impl/Kokkos_Error.hpp>
28#include <complex>
29#include <type_traits>
30#include <iosfwd>
31
32namespace Kokkos {
33
41template <class RealType>
42class
43#ifdef KOKKOS_ENABLE_COMPLEX_ALIGN
44 alignas(2 * sizeof(RealType))
45#endif
46 complex {
47 private:
48 RealType re_{};
49 RealType im_{};
50
51 public:
53 using value_type = RealType;
54
56 KOKKOS_DEFAULTED_FUNCTION
57 complex() = default;
58
60 KOKKOS_DEFAULTED_FUNCTION
61 complex(const complex&) noexcept = default;
62
63 KOKKOS_DEFAULTED_FUNCTION
64 complex& operator=(const complex&) noexcept = default;
65
67 template <
68 class RType,
69 std::enable_if_t<std::is_convertible<RType, RealType>::value, int> = 0>
70 KOKKOS_INLINE_FUNCTION complex(const complex<RType>& other) noexcept
71 // Intentionally do the conversions implicitly here so that users don't
72 // get any warnings about narrowing, etc., that they would expect to get
73 // otherwise.
74 : re_(other.real()), im_(other.imag()) {}
75
81 KOKKOS_INLINE_FUNCTION
82 complex(const std::complex<RealType>& src) noexcept
83 // We can use this aspect of the standard to avoid calling
84 // non-device-marked functions `std::real` and `std::imag`: "For any
85 // object z of type complex<T>, reinterpret_cast<T(&)[2]>(z)[0] is the
86 // real part of z and reinterpret_cast<T(&)[2]>(z)[1] is the imaginary
87 // part of z." Now we don't have to provide a whole bunch of the overloads
88 // of things taking either Kokkos::complex or std::complex
89 : re_(reinterpret_cast<const RealType (&)[2]>(src)[0]),
90 im_(reinterpret_cast<const RealType (&)[2]>(src)[1]) {}
91
97 // TODO: make explicit. DJS 2019-08-28
98 operator std::complex<RealType>() const noexcept {
99 return std::complex<RealType>(re_, im_);
100 }
101
104 KOKKOS_INLINE_FUNCTION complex(const RealType& val) noexcept
105 : re_(val), im_(static_cast<RealType>(0)) {}
106
108 KOKKOS_INLINE_FUNCTION
109 complex(const RealType& re, const RealType& im) noexcept : re_(re), im_(im) {}
110
112 KOKKOS_INLINE_FUNCTION complex& operator=(const RealType& val) noexcept {
113 re_ = val;
114 im_ = RealType(0);
115 return *this;
116 }
117
123 complex& operator=(const std::complex<RealType>& src) noexcept {
124 *this = complex(src);
125 return *this;
126 }
127
129 KOKKOS_INLINE_FUNCTION
130 constexpr RealType& imag() noexcept { return im_; }
131
133 KOKKOS_INLINE_FUNCTION
134 constexpr RealType& real() noexcept { return re_; }
135
137 KOKKOS_INLINE_FUNCTION
138 constexpr RealType imag() const noexcept { return im_; }
139
141 KOKKOS_INLINE_FUNCTION
142 constexpr RealType real() const noexcept { return re_; }
143
145 KOKKOS_INLINE_FUNCTION
146 constexpr void imag(RealType v) noexcept { im_ = v; }
147
149 KOKKOS_INLINE_FUNCTION
150 constexpr void real(RealType v) noexcept { re_ = v; }
151
152 constexpr KOKKOS_INLINE_FUNCTION complex& operator+=(
153 const complex<RealType>& src) noexcept {
154 re_ += src.re_;
155 im_ += src.im_;
156 return *this;
157 }
158
159 constexpr KOKKOS_INLINE_FUNCTION complex& operator+=(
160 const RealType& src) noexcept {
161 re_ += src;
162 return *this;
163 }
164
165 constexpr KOKKOS_INLINE_FUNCTION complex& operator-=(
166 const complex<RealType>& src) noexcept {
167 re_ -= src.re_;
168 im_ -= src.im_;
169 return *this;
170 }
171
172 constexpr KOKKOS_INLINE_FUNCTION complex& operator-=(
173 const RealType& src) noexcept {
174 re_ -= src;
175 return *this;
176 }
177
178 constexpr KOKKOS_INLINE_FUNCTION complex& operator*=(
179 const complex<RealType>& src) noexcept {
180 const RealType realPart = re_ * src.re_ - im_ * src.im_;
181 const RealType imagPart = re_ * src.im_ + im_ * src.re_;
182 re_ = realPart;
183 im_ = imagPart;
184 return *this;
185 }
186
187 constexpr KOKKOS_INLINE_FUNCTION complex& operator*=(
188 const RealType& src) noexcept {
189 re_ *= src;
190 im_ *= src;
191 return *this;
192 }
193
194 // Conditional noexcept, just in case RType throws on divide-by-zero
195 constexpr KOKKOS_INLINE_FUNCTION complex& operator/=(
196 const complex<RealType>& y) noexcept(noexcept(RealType{} / RealType{})) {
197 // Scale (by the "1-norm" of y) to avoid unwarranted overflow.
198 // If the real part is +/-Inf and the imaginary part is -/+Inf,
199 // this won't change the result.
200 const RealType s = fabs(y.real()) + fabs(y.imag());
201
202 // If s is 0, then y is zero, so x/y == real(x)/0 + i*imag(x)/0.
203 // In that case, the relation x/y == (x/s) / (y/s) doesn't hold,
204 // because y/s is NaN.
205 // TODO mark this branch unlikely
206 if (s == RealType(0)) {
207 this->re_ /= s;
208 this->im_ /= s;
209 } else {
210 const complex x_scaled(this->re_ / s, this->im_ / s);
211 const complex y_conj_scaled(y.re_ / s, -(y.im_) / s);
212 const RealType y_scaled_abs =
213 y_conj_scaled.re_ * y_conj_scaled.re_ +
214 y_conj_scaled.im_ * y_conj_scaled.im_; // abs(y) == abs(conj(y))
215 *this = x_scaled * y_conj_scaled;
216 *this /= y_scaled_abs;
217 }
218 return *this;
219 }
220
221 constexpr KOKKOS_INLINE_FUNCTION complex& operator/=(
222 const std::complex<RealType>& y) noexcept(noexcept(RealType{} /
223 RealType{})) {
224 // Scale (by the "1-norm" of y) to avoid unwarranted overflow.
225 // If the real part is +/-Inf and the imaginary part is -/+Inf,
226 // this won't change the result.
227 const RealType s = fabs(y.real()) + fabs(y.imag());
228
229 // If s is 0, then y is zero, so x/y == real(x)/0 + i*imag(x)/0.
230 // In that case, the relation x/y == (x/s) / (y/s) doesn't hold,
231 // because y/s is NaN.
232 if (s == RealType(0)) {
233 this->re_ /= s;
234 this->im_ /= s;
235 } else {
236 const complex x_scaled(this->re_ / s, this->im_ / s);
237 const complex y_conj_scaled(y.re_ / s, -(y.im_) / s);
238 const RealType y_scaled_abs =
239 y_conj_scaled.re_ * y_conj_scaled.re_ +
240 y_conj_scaled.im_ * y_conj_scaled.im_; // abs(y) == abs(conj(y))
241 *this = x_scaled * y_conj_scaled;
242 *this /= y_scaled_abs;
243 }
244 return *this;
245 }
246
247 constexpr KOKKOS_INLINE_FUNCTION complex& operator/=(
248 const RealType& src) noexcept(noexcept(RealType{} / RealType{})) {
249 re_ /= src;
250 im_ /= src;
251 return *this;
252 }
253
254#ifdef KOKKOS_ENABLE_DEPRECATED_CODE_4
256 template <
257 class RType,
258 std::enable_if_t<std::is_convertible<RType, RealType>::value, int> = 0>
259 KOKKOS_DEPRECATED KOKKOS_INLINE_FUNCTION
260 complex(const volatile complex<RType>& src) noexcept
261 // Intentionally do the conversions implicitly here so that users don't
262 // get any warnings about narrowing, etc., that they would expect to get
263 // otherwise.
264 : re_(src.re_), im_(src.im_) {}
265
275 //
276 // Templated, so as not to be a copy assignment operator (Kokkos issue #2577)
277 // Intended to behave as
278 // void operator=(const complex&) volatile noexcept
279 //
280 // Use cases:
281 // complex r;
282 // const complex cr;
283 // volatile complex vl;
284 // vl = r;
285 // vl = cr;
286 template <class Complex,
287 std::enable_if_t<std::is_same<Complex, complex>::value, int> = 0>
288 KOKKOS_DEPRECATED KOKKOS_INLINE_FUNCTION void operator=(
289 const Complex& src) volatile noexcept {
290 re_ = src.re_;
291 im_ = src.im_;
292 // We deliberately do not return anything here. See explanation
293 // in public documentation above.
294 }
295
297 // TODO Should this return void like the other volatile assignment operators?
298 //
299 // Templated, so as not to be a copy assignment operator (Kokkos issue #2577)
300 // Intended to behave as
301 // volatile complex& operator=(const volatile complex&) volatile noexcept
302 //
303 // Use cases:
304 // volatile complex vr;
305 // const volatile complex cvr;
306 // volatile complex vl;
307 // vl = vr;
308 // vl = cvr;
309 template <class Complex,
310 std::enable_if_t<std::is_same<Complex, complex>::value, int> = 0>
311 KOKKOS_DEPRECATED KOKKOS_INLINE_FUNCTION volatile complex& operator=(
312 const volatile Complex& src) volatile noexcept {
313 re_ = src.re_;
314 im_ = src.im_;
315 return *this;
316 }
317
319 //
320 // Templated, so as not to be a copy assignment operator (Kokkos issue #2577)
321 // Intended to behave as
322 // complex& operator=(const volatile complex&) noexcept
323 //
324 // Use cases:
325 // volatile complex vr;
326 // const volatile complex cvr;
327 // complex l;
328 // l = vr;
329 // l = cvr;
330 //
331 template <class Complex,
332 std::enable_if_t<std::is_same<Complex, complex>::value, int> = 0>
333 KOKKOS_DEPRECATED KOKKOS_INLINE_FUNCTION complex& operator=(
334 const volatile Complex& src) noexcept {
335 re_ = src.re_;
336 im_ = src.im_;
337 return *this;
338 }
339
340 // Mirroring the behavior of the assignment operators from complex RHS in the
341 // RealType RHS versions.
342
344 KOKKOS_DEPRECATED KOKKOS_INLINE_FUNCTION void operator=(
345 const volatile RealType& val) noexcept {
346 re_ = val;
347 im_ = RealType(0);
348 // We deliberately do not return anything here. See explanation
349 // in public documentation above.
350 }
351
353 KOKKOS_DEPRECATED KOKKOS_INLINE_FUNCTION complex& operator=(
354 const RealType& val) volatile noexcept {
355 re_ = val;
356 im_ = RealType(0);
357 return *this;
358 }
359
361 // TODO Should this return void like the other volatile assignment operators?
362 KOKKOS_DEPRECATED KOKKOS_INLINE_FUNCTION complex& operator=(
363 const volatile RealType& val) volatile noexcept {
364 re_ = val;
365 im_ = RealType(0);
366 return *this;
367 }
368
370 KOKKOS_DEPRECATED KOKKOS_INLINE_FUNCTION volatile RealType&
371 imag() volatile noexcept {
372 return im_;
373 }
374
376 KOKKOS_DEPRECATED KOKKOS_INLINE_FUNCTION volatile RealType&
377 real() volatile noexcept {
378 return re_;
379 }
380
382 KOKKOS_DEPRECATED KOKKOS_INLINE_FUNCTION RealType imag() const
383 volatile noexcept {
384 return im_;
385 }
386
388 KOKKOS_DEPRECATED KOKKOS_INLINE_FUNCTION RealType real() const
389 volatile noexcept {
390 return re_;
391 }
392
393 KOKKOS_DEPRECATED KOKKOS_INLINE_FUNCTION void operator+=(
394 const volatile complex<RealType>& src) volatile noexcept {
395 re_ += src.re_;
396 im_ += src.im_;
397 }
398
399 KOKKOS_DEPRECATED KOKKOS_INLINE_FUNCTION void operator+=(
400 const volatile RealType& src) volatile noexcept {
401 re_ += src;
402 }
403
404 KOKKOS_DEPRECATED KOKKOS_INLINE_FUNCTION void operator*=(
405 const volatile complex<RealType>& src) volatile noexcept {
406 const RealType realPart = re_ * src.re_ - im_ * src.im_;
407 const RealType imagPart = re_ * src.im_ + im_ * src.re_;
408
409 re_ = realPart;
410 im_ = imagPart;
411 }
412
413 KOKKOS_DEPRECATED KOKKOS_INLINE_FUNCTION void operator*=(
414 const volatile RealType& src) volatile noexcept {
415 re_ *= src;
416 im_ *= src;
417 }
418#endif // KOKKOS_ENABLE_DEPRECATED_CODE_4
419};
420
421//==============================================================================
422// <editor-fold desc="Equality and inequality"> {{{1
423
424// Note that this is not the same behavior as std::complex, which doesn't allow
425// implicit conversions, but since this is the way we had it before, we have
426// to do it this way now.
427
429template <class RealType1, class RealType2>
430KOKKOS_INLINE_FUNCTION bool operator==(complex<RealType1> const& x,
431 complex<RealType2> const& y) noexcept {
432 using common_type = std::common_type_t<RealType1, RealType2>;
433 return common_type(x.real()) == common_type(y.real()) &&
434 common_type(x.imag()) == common_type(y.imag());
435}
436
437// TODO (here and elsewhere) decide if we should convert to a Kokkos::complex
438// and do the comparison in a device-marked function
440template <class RealType1, class RealType2>
441inline bool operator==(std::complex<RealType1> const& x,
442 complex<RealType2> const& y) noexcept {
443 using common_type = std::common_type_t<RealType1, RealType2>;
444 return common_type(x.real()) == common_type(y.real()) &&
445 common_type(x.imag()) == common_type(y.imag());
446}
447
449template <class RealType1, class RealType2>
450inline bool operator==(complex<RealType1> const& x,
451 std::complex<RealType2> const& y) noexcept {
452 using common_type = std::common_type_t<RealType1, RealType2>;
453 return common_type(x.real()) == common_type(y.real()) &&
454 common_type(x.imag()) == common_type(y.imag());
455}
456
458template <
459 class RealType1, class RealType2,
460 // Constraints to avoid participation in oparator==() for every possible RHS
461 std::enable_if_t<std::is_convertible<RealType2, RealType1>::value, int> = 0>
462KOKKOS_INLINE_FUNCTION bool operator==(complex<RealType1> const& x,
463 RealType2 const& y) noexcept {
464 using common_type = std::common_type_t<RealType1, RealType2>;
465 return common_type(x.real()) == common_type(y) &&
466 common_type(x.imag()) == common_type(0);
467}
468
470template <
471 class RealType1, class RealType2,
472 // Constraints to avoid participation in oparator==() for every possible RHS
473 std::enable_if_t<std::is_convertible<RealType1, RealType2>::value, int> = 0>
474KOKKOS_INLINE_FUNCTION bool operator==(RealType1 const& x,
475 complex<RealType2> const& y) noexcept {
476 using common_type = std::common_type_t<RealType1, RealType2>;
477 return common_type(x) == common_type(y.real()) &&
478 common_type(0) == common_type(y.imag());
479}
480
482template <class RealType1, class RealType2>
483KOKKOS_INLINE_FUNCTION bool operator!=(complex<RealType1> const& x,
484 complex<RealType2> const& y) noexcept {
485 using common_type = std::common_type_t<RealType1, RealType2>;
486 return common_type(x.real()) != common_type(y.real()) ||
487 common_type(x.imag()) != common_type(y.imag());
488}
489
491template <class RealType1, class RealType2>
492inline bool operator!=(std::complex<RealType1> const& x,
493 complex<RealType2> const& y) noexcept {
494 using common_type = std::common_type_t<RealType1, RealType2>;
495 return common_type(x.real()) != common_type(y.real()) ||
496 common_type(x.imag()) != common_type(y.imag());
497}
498
500template <class RealType1, class RealType2>
501inline bool operator!=(complex<RealType1> const& x,
502 std::complex<RealType2> const& y) noexcept {
503 using common_type = std::common_type_t<RealType1, RealType2>;
504 return common_type(x.real()) != common_type(y.real()) ||
505 common_type(x.imag()) != common_type(y.imag());
506}
507
509template <
510 class RealType1, class RealType2,
511 // Constraints to avoid participation in oparator==() for every possible RHS
512 std::enable_if_t<std::is_convertible<RealType2, RealType1>::value, int> = 0>
513KOKKOS_INLINE_FUNCTION bool operator!=(complex<RealType1> const& x,
514 RealType2 const& y) noexcept {
515 using common_type = std::common_type_t<RealType1, RealType2>;
516 return common_type(x.real()) != common_type(y) ||
517 common_type(x.imag()) != common_type(0);
518}
519
521template <
522 class RealType1, class RealType2,
523 // Constraints to avoid participation in oparator==() for every possible RHS
524 std::enable_if_t<std::is_convertible<RealType1, RealType2>::value, int> = 0>
525KOKKOS_INLINE_FUNCTION bool operator!=(RealType1 const& x,
526 complex<RealType2> const& y) noexcept {
527 using common_type = std::common_type_t<RealType1, RealType2>;
528 return common_type(x) != common_type(y.real()) ||
529 common_type(0) != common_type(y.imag());
530}
531
532// </editor-fold> end Equality and inequality }}}1
533//==============================================================================
534
536template <class RealType1, class RealType2>
537KOKKOS_INLINE_FUNCTION complex<std::common_type_t<RealType1, RealType2>>
538operator+(const complex<RealType1>& x, const complex<RealType2>& y) noexcept {
539 return complex<std::common_type_t<RealType1, RealType2>>(x.real() + y.real(),
540 x.imag() + y.imag());
541}
542
544template <class RealType1, class RealType2>
545KOKKOS_INLINE_FUNCTION complex<std::common_type_t<RealType1, RealType2>>
546operator+(const complex<RealType1>& x, const RealType2& y) noexcept {
547 return complex<std::common_type_t<RealType1, RealType2>>(x.real() + y,
548 x.imag());
549}
550
552template <class RealType1, class RealType2>
553KOKKOS_INLINE_FUNCTION complex<std::common_type_t<RealType1, RealType2>>
554operator+(const RealType1& x, const complex<RealType2>& y) noexcept {
555 return complex<std::common_type_t<RealType1, RealType2>>(x + y.real(),
556 y.imag());
557}
558
560template <class RealType>
561KOKKOS_INLINE_FUNCTION complex<RealType> operator+(
562 const complex<RealType>& x) noexcept {
563 return complex<RealType>{+x.real(), +x.imag()};
564}
565
567template <class RealType1, class RealType2>
568KOKKOS_INLINE_FUNCTION complex<std::common_type_t<RealType1, RealType2>>
569operator-(const complex<RealType1>& x, const complex<RealType2>& y) noexcept {
570 return complex<std::common_type_t<RealType1, RealType2>>(x.real() - y.real(),
571 x.imag() - y.imag());
572}
573
575template <class RealType1, class RealType2>
576KOKKOS_INLINE_FUNCTION complex<std::common_type_t<RealType1, RealType2>>
577operator-(const complex<RealType1>& x, const RealType2& y) noexcept {
578 return complex<std::common_type_t<RealType1, RealType2>>(x.real() - y,
579 x.imag());
580}
581
583template <class RealType1, class RealType2>
584KOKKOS_INLINE_FUNCTION complex<std::common_type_t<RealType1, RealType2>>
585operator-(const RealType1& x, const complex<RealType2>& y) noexcept {
586 return complex<std::common_type_t<RealType1, RealType2>>(x - y.real(),
587 -y.imag());
588}
589
591template <class RealType>
592KOKKOS_INLINE_FUNCTION complex<RealType> operator-(
593 const complex<RealType>& x) noexcept {
594 return complex<RealType>(-x.real(), -x.imag());
595}
596
598template <class RealType1, class RealType2>
599KOKKOS_INLINE_FUNCTION complex<std::common_type_t<RealType1, RealType2>>
600operator*(const complex<RealType1>& x, const complex<RealType2>& y) noexcept {
601 return complex<std::common_type_t<RealType1, RealType2>>(
602 x.real() * y.real() - x.imag() * y.imag(),
603 x.real() * y.imag() + x.imag() * y.real());
604}
605
614template <class RealType1, class RealType2>
615inline complex<std::common_type_t<RealType1, RealType2>> operator*(
616 const std::complex<RealType1>& x, const complex<RealType2>& y) {
617 return complex<std::common_type_t<RealType1, RealType2>>(
618 x.real() * y.real() - x.imag() * y.imag(),
619 x.real() * y.imag() + x.imag() * y.real());
620}
621
626template <class RealType1, class RealType2>
627KOKKOS_INLINE_FUNCTION complex<std::common_type_t<RealType1, RealType2>>
628operator*(const RealType1& x, const complex<RealType2>& y) noexcept {
629 return complex<std::common_type_t<RealType1, RealType2>>(x * y.real(),
630 x * y.imag());
631}
632
637template <class RealType1, class RealType2>
638KOKKOS_INLINE_FUNCTION complex<std::common_type_t<RealType1, RealType2>>
639operator*(const complex<RealType1>& y, const RealType2& x) noexcept {
640 return complex<std::common_type_t<RealType1, RealType2>>(x * y.real(),
641 x * y.imag());
642}
643
645template <class RealType>
646KOKKOS_INLINE_FUNCTION RealType imag(const complex<RealType>& x) noexcept {
647 return x.imag();
648}
649
650template <class ArithmeticType>
651KOKKOS_INLINE_FUNCTION constexpr Impl::promote_t<ArithmeticType> imag(
652 ArithmeticType) {
653 return ArithmeticType();
654}
655
657template <class RealType>
658KOKKOS_INLINE_FUNCTION RealType real(const complex<RealType>& x) noexcept {
659 return x.real();
660}
661
662template <class ArithmeticType>
663KOKKOS_INLINE_FUNCTION constexpr Impl::promote_t<ArithmeticType> real(
664 ArithmeticType x) {
665 return x;
666}
667
669template <class T>
670KOKKOS_INLINE_FUNCTION complex<T> polar(const T& r, const T& theta = T()) {
671 KOKKOS_EXPECTS(r >= 0);
672 return complex<T>(r * cos(theta), r * sin(theta));
673}
674
676template <class RealType>
677KOKKOS_INLINE_FUNCTION RealType abs(const complex<RealType>& x) {
678 return hypot(x.real(), x.imag());
679}
680
682template <class T>
683KOKKOS_INLINE_FUNCTION complex<T> pow(const complex<T>& x, const T& y) {
684 T r = abs(x);
685 T theta = atan2(x.imag(), x.real());
686 return polar(pow(r, y), y * theta);
687}
688
689template <class T>
690KOKKOS_INLINE_FUNCTION complex<T> pow(const T& x, const complex<T>& y) {
691 return pow(complex<T>(x), y);
692}
693
694template <class T>
695KOKKOS_INLINE_FUNCTION complex<T> pow(const complex<T>& x,
696 const complex<T>& y) {
697 return x == T() ? T() : exp(y * log(x));
698}
699
700template <class T, class U,
701 class = std::enable_if_t<std::is_arithmetic<T>::value>>
702KOKKOS_INLINE_FUNCTION complex<Impl::promote_2_t<T, U>> pow(
703 const T& x, const complex<U>& y) {
704 using type = Impl::promote_2_t<T, U>;
705 return pow(type(x), complex<type>(y));
706}
707
708template <class T, class U,
709 class = std::enable_if_t<std::is_arithmetic<U>::value>>
710KOKKOS_INLINE_FUNCTION complex<Impl::promote_2_t<T, U>> pow(const complex<T>& x,
711 const U& y) {
712 using type = Impl::promote_2_t<T, U>;
713 return pow(complex<type>(x), type(y));
714}
715
716template <class T, class U>
717KOKKOS_INLINE_FUNCTION complex<Impl::promote_2_t<T, U>> pow(
718 const complex<T>& x, const complex<U>& y) {
719 using type = Impl::promote_2_t<T, U>;
720 return pow(complex<type>(x), complex<type>(y));
721}
722
725template <class RealType>
726KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> sqrt(
727 const complex<RealType>& x) {
728 RealType r = x.real();
729 RealType i = x.imag();
730
731 if (r == RealType()) {
732 RealType t = sqrt(fabs(i) / 2);
733 return Kokkos::complex<RealType>(t, i < RealType() ? -t : t);
734 } else {
735 RealType t = sqrt(2 * (abs(x) + fabs(r)));
736 RealType u = t / 2;
737 return r > RealType() ? Kokkos::complex<RealType>(u, i / t)
738 : Kokkos::complex<RealType>(fabs(i) / t,
739 i < RealType() ? -u : u);
740 }
741}
742
744template <class RealType>
745KOKKOS_INLINE_FUNCTION complex<RealType> conj(
746 const complex<RealType>& x) noexcept {
747 return complex<RealType>(real(x), -imag(x));
748}
749
750template <class ArithmeticType>
751KOKKOS_INLINE_FUNCTION constexpr complex<Impl::promote_t<ArithmeticType>> conj(
752 ArithmeticType x) {
753 using type = Impl::promote_t<ArithmeticType>;
754 return complex<type>(x, -type());
755}
756
758template <class RealType>
759KOKKOS_INLINE_FUNCTION complex<RealType> exp(const complex<RealType>& x) {
760 return exp(x.real()) * complex<RealType>(cos(x.imag()), sin(x.imag()));
761}
762
764template <class RealType>
765KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> log(
766 const complex<RealType>& x) {
767 RealType phi = atan2(x.imag(), x.real());
768 return Kokkos::complex<RealType>(log(abs(x)), phi);
769}
770
772template <class RealType>
773KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> log10(
774 const complex<RealType>& x) {
775 return log(x) / log(RealType(10));
776}
777
779template <class RealType>
780KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> sin(
781 const complex<RealType>& x) {
782 return Kokkos::complex<RealType>(sin(x.real()) * cosh(x.imag()),
783 cos(x.real()) * sinh(x.imag()));
784}
785
787template <class RealType>
788KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> cos(
789 const complex<RealType>& x) {
790 return Kokkos::complex<RealType>(cos(x.real()) * cosh(x.imag()),
791 -sin(x.real()) * sinh(x.imag()));
792}
793
795template <class RealType>
796KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> tan(
797 const complex<RealType>& x) {
798 return sin(x) / cos(x);
799}
800
802template <class RealType>
803KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> sinh(
804 const complex<RealType>& x) {
805 return Kokkos::complex<RealType>(sinh(x.real()) * cos(x.imag()),
806 cosh(x.real()) * sin(x.imag()));
807}
808
810template <class RealType>
811KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> cosh(
812 const complex<RealType>& x) {
813 return Kokkos::complex<RealType>(cosh(x.real()) * cos(x.imag()),
814 sinh(x.real()) * sin(x.imag()));
815}
816
818template <class RealType>
819KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> tanh(
820 const complex<RealType>& x) {
821 return sinh(x) / cosh(x);
822}
823
825template <class RealType>
826KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> asinh(
827 const complex<RealType>& x) {
828 return log(x + sqrt(x * x + RealType(1.0)));
829}
830
832template <class RealType>
833KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> acosh(
834 const complex<RealType>& x) {
835 return RealType(2.0) * log(sqrt(RealType(0.5) * (x + RealType(1.0))) +
836 sqrt(RealType(0.5) * (x - RealType(1.0))));
837}
838
840template <class RealType>
841KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> atanh(
842 const complex<RealType>& x) {
843 const RealType i2 = x.imag() * x.imag();
844 const RealType r = RealType(1.0) - i2 - x.real() * x.real();
845
846 RealType p = RealType(1.0) + x.real();
847 RealType m = RealType(1.0) - x.real();
848
849 p = i2 + p * p;
850 m = i2 + m * m;
851
852 RealType phi = atan2(RealType(2.0) * x.imag(), r);
853 return Kokkos::complex<RealType>(RealType(0.25) * (log(p) - log(m)),
854 RealType(0.5) * phi);
855}
856
858template <class RealType>
859KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> asin(
860 const complex<RealType>& x) {
862 asinh(Kokkos::complex<RealType>(-x.imag(), x.real()));
863 return Kokkos::complex<RealType>(t.imag(), -t.real());
864}
865
867template <class RealType>
868KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> acos(
869 const complex<RealType>& x) {
870 Kokkos::complex<RealType> t = asin(x);
871 RealType pi_2 = acos(RealType(0.0));
872 return Kokkos::complex<RealType>(pi_2 - t.real(), -t.imag());
873}
874
876template <class RealType>
877KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> atan(
878 const complex<RealType>& x) {
879 const RealType r2 = x.real() * x.real();
880 const RealType i = RealType(1.0) - r2 - x.imag() * x.imag();
881
882 RealType p = x.imag() + RealType(1.0);
883 RealType m = x.imag() - RealType(1.0);
884
885 p = r2 + p * p;
886 m = r2 + m * m;
887
889 RealType(0.5) * atan2(RealType(2.0) * x.real(), i),
890 RealType(0.25) * log(p / m));
891}
892
896template <class RealType>
897inline complex<RealType> exp(const std::complex<RealType>& c) {
898 return complex<RealType>(std::exp(c.real()) * std::cos(c.imag()),
899 std::exp(c.real()) * std::sin(c.imag()));
900}
901
903template <class RealType1, class RealType2>
904KOKKOS_INLINE_FUNCTION complex<std::common_type_t<RealType1, RealType2>>
905operator/(const complex<RealType1>& x,
906 const RealType2& y) noexcept(noexcept(RealType1{} / RealType2{})) {
907 return complex<std::common_type_t<RealType1, RealType2>>(real(x) / y,
908 imag(x) / y);
909}
910
912template <class RealType1, class RealType2>
913KOKKOS_INLINE_FUNCTION complex<std::common_type_t<RealType1, RealType2>>
914operator/(const complex<RealType1>& x,
915 const complex<RealType2>& y) noexcept(noexcept(RealType1{} /
916 RealType2{})) {
917 // Scale (by the "1-norm" of y) to avoid unwarranted overflow.
918 // If the real part is +/-Inf and the imaginary part is -/+Inf,
919 // this won't change the result.
920 using common_real_type = std::common_type_t<RealType1, RealType2>;
921 const common_real_type s = fabs(real(y)) + fabs(imag(y));
922
923 // If s is 0, then y is zero, so x/y == real(x)/0 + i*imag(x)/0.
924 // In that case, the relation x/y == (x/s) / (y/s) doesn't hold,
925 // because y/s is NaN.
926 if (s == 0.0) {
927 return complex<common_real_type>(real(x) / s, imag(x) / s);
928 } else {
929 const complex<common_real_type> x_scaled(real(x) / s, imag(x) / s);
930 const complex<common_real_type> y_conj_scaled(real(y) / s, -imag(y) / s);
931 const RealType1 y_scaled_abs =
932 real(y_conj_scaled) * real(y_conj_scaled) +
933 imag(y_conj_scaled) * imag(y_conj_scaled); // abs(y) == abs(conj(y))
934 complex<common_real_type> result = x_scaled * y_conj_scaled;
935 result /= y_scaled_abs;
936 return result;
937 }
938}
939
941template <class RealType1, class RealType2>
942KOKKOS_INLINE_FUNCTION complex<std::common_type_t<RealType1, RealType2>>
943operator/(const RealType1& x,
944 const complex<RealType2>& y) noexcept(noexcept(RealType1{} /
945 RealType2{})) {
946 return complex<std::common_type_t<RealType1, RealType2>>(x) / y;
947}
948
949template <class RealType>
950std::ostream& operator<<(std::ostream& os, const complex<RealType>& x) {
951 const std::complex<RealType> x_std(Kokkos::real(x), Kokkos::imag(x));
952 os << x_std;
953 return os;
954}
955
956template <class RealType>
957std::istream& operator>>(std::istream& is, complex<RealType>& x) {
958 std::complex<RealType> x_std;
959 is >> x_std;
960 x = x_std; // only assigns on success of above
961 return is;
962}
963
964template <class T>
965struct reduction_identity<Kokkos::complex<T>> {
966 using t_red_ident = reduction_identity<T>;
967 KOKKOS_FORCEINLINE_FUNCTION constexpr static Kokkos::complex<T>
968 sum() noexcept {
969 return Kokkos::complex<T>(t_red_ident::sum(), t_red_ident::sum());
970 }
971 KOKKOS_FORCEINLINE_FUNCTION constexpr static Kokkos::complex<T>
972 prod() noexcept {
973 return Kokkos::complex<T>(t_red_ident::prod(), t_red_ident::sum());
974 }
975};
976
977} // namespace Kokkos
978
979#ifdef KOKKOS_IMPL_PUBLIC_INCLUDE_NOTDEFINED_COMPLEX
980#undef KOKKOS_IMPL_PUBLIC_INCLUDE
981#undef KOKKOS_IMPL_PUBLIC_INCLUDE_NOTDEFINED_COMPLEX
982#endif
983#endif // KOKKOS_COMPLEX_HPP
Atomic functions.
Partial reimplementation of std::complex that works as the result of a Kokkos::parallel_reduce.
KOKKOS_INLINE_FUNCTION constexpr RealType imag() const noexcept
The imaginary part of this complex number.
RealType value_type
The type of the real or imaginary parts of this complex number.
KOKKOS_INLINE_FUNCTION constexpr RealType & real() noexcept
The real part of this complex number.
complex & operator=(const std::complex< RealType > &src) noexcept
Assignment operator from std::complex.
KOKKOS_DEFAULTED_FUNCTION complex(const complex &) noexcept=default
Copy constructor.
KOKKOS_INLINE_FUNCTION constexpr void real(RealType v) noexcept
Set the real part of this complex number.
KOKKOS_INLINE_FUNCTION complex(const RealType &re, const RealType &im) noexcept
Constructor that takes the real and imaginary parts.
KOKKOS_INLINE_FUNCTION constexpr RealType & imag() noexcept
The imaginary part of this complex number.
KOKKOS_INLINE_FUNCTION constexpr RealType real() const noexcept
The real part of this complex number.
KOKKOS_INLINE_FUNCTION complex & operator=(const RealType &val) noexcept
Assignment operator (from a real number).
KOKKOS_INLINE_FUNCTION constexpr void imag(RealType v) noexcept
Set the imaginary part of this complex number.
KOKKOS_DEFAULTED_FUNCTION complex()=default
Default constructor (initializes both real and imaginary parts to zero).
KOKKOS_INLINE_FUNCTION complex(const RealType &val) noexcept
Constructor that takes just the real part, and sets the imaginary part to zero.
KOKKOS_INLINE_FUNCTION complex(const std::complex< RealType > &src) noexcept
Conversion constructor from std::complex.
KOKKOS_INLINE_FUNCTION complex(const complex< RType > &other) noexcept
Conversion constructor from compatible RType.