prdc/cuda/complex.h

#ifndef PRDC_CUDA_COMPLEX_H
#define PRDC_CUDA_COMPLEX_H
 
/*
* PSCF Package - Polymer Self-Consistent Field 
*
* Copyright 2016 - 2023, The Regents of the University of Minnesota
* Distributed under the terms of the GNU General Public License.
*/
 
#include "types.h"
#include <pscf/math/complex.h>
#include <complex>
 
namespace Pscf {
 
   /*
   * Types cudaComplex and cudaReal are defined in prdc/cuda/types.h
   * (in the Prdc::Cuda namespace) as aliases for cufft complex and 
   * real types. They may be either single or double precision.
   */
 
   // Real and imaginary components
 
   template <> inline 
   Prdc::Cuda::cudaReal real(Prdc::Cuda::cudaComplex const& a)
   {  return a.x; }
 
   template <> inline 
   Prdc::Cuda::cudaReal imag(Prdc::Cuda::cudaComplex const& a)
   {  return a.y; }
 
   // Absolute magnitude
 
   template <> inline 
   Prdc::Cuda::cudaReal abs(Prdc::Cuda::cudaComplex const& a)
   {  return sqrt(a.x * a.x + a.y * a.y); }
 
   template <> inline 
   Prdc::Cuda::cudaReal absSq(Prdc::Cuda::cudaComplex const& a)
   {  return (a.x * a.x + a.y * a.y); }
 
   // cudaComplex Conjugation
 
   template <> inline
   void conj(Prdc::Cuda::cudaComplex& z, Prdc::Cuda::cudaComplex const& a)
   {
      z.x = a.x;
      z.y = -a.y;
   }
 
   template <> inline
   void conj(Prdc::Cuda::cudaComplex& a)
   {
      a.x = a.x;
      a.y = -a.y;
   }
 
   // Assignment
 
   template <> inline
   void assign(Prdc::Cuda::cudaComplex& z, Prdc::Cuda::cudaReal const& a, 
               Prdc::Cuda::cudaReal const& b)
   {
      z.x = a;
      z.y = b;
   }
 
   template <> inline
   void assign(Prdc::Cuda::cudaComplex& z, Prdc::Cuda::cudaReal const& a)
   {
      z.x = a;
      z.y = 0;
   }
 
   template <> inline
   void assign(Prdc::Cuda::cudaComplex& z, Prdc::Cuda::cudaComplex const& a)
   {
      z.x = a.x;
      z.y = a.y;
   }
 
   template <> inline
   void assign(Prdc::Cuda::cudaComplex & z, 
               std::complex<Prdc::Cuda::cudaReal> const& a)
   {
      z.x = a.real();
      z.y = a.imag();
   }
 
   template <> inline
   void assign(std::complex<Prdc::Cuda::cudaReal> & z, 
               Prdc::Cuda::cudaComplex const& a)
   {  z = std::complex<Prdc::Cuda::cudaReal>(a.x, a.y); }
 
   // Addition
 
   template <> inline
   void add(Prdc::Cuda::cudaComplex& z, Prdc::Cuda::cudaComplex const& a, 
            Prdc::Cuda::cudaComplex const& b)
   {
      z.x = a.x + b.x;
      z.y = a.y + b.y;
   }
 
   template <> inline
   void add(Prdc::Cuda::cudaComplex& z, Prdc::Cuda::cudaComplex const& a, 
            Prdc::Cuda::cudaReal const& b)
   {
      z.x = a.x + b;
      z.y = a.y;
   }
 
   template <> inline
   void addEq(Prdc::Cuda::cudaComplex& a, Prdc::Cuda::cudaComplex const& b)
   {
      a.x += b.x;
      a.y += b.y;
   }
 
   template <> inline
   void addEq(Prdc::Cuda::cudaComplex& a, Prdc::Cuda::cudaReal const& b)
   {
      a.x += b;
   }
 
   // Subtraction
 
   template <> inline
   void sub(Prdc::Cuda::cudaComplex& z, Prdc::Cuda::cudaComplex const& a, 
            Prdc::Cuda::cudaComplex const& b)
   {
      z.x = a.x - b.x;
      z.y = a.y - b.y;
   }
 
   template <> inline
   void sub(Prdc::Cuda::cudaComplex& z, Prdc::Cuda::cudaComplex const& a, 
            Prdc::Cuda::cudaReal const& b)
   {
      z.x = a.x - b;
      z.y = a.y;
   }
 
   template <> inline
   void subEq(Prdc::Cuda::cudaComplex & a, Prdc::Cuda::cudaComplex const& b)
   {
      a.x -= b.x;
      a.y -= b.y;
   }
 
   template <> inline
   void subEq(Prdc::Cuda::cudaComplex & a, Prdc::Cuda::cudaReal const& b)
   {
      a.x -= b;
   }
 
   template <> inline 
   Prdc::Cuda::cudaReal absSqDiff(Prdc::Cuda::cudaComplex const& a, 
                                  Prdc::Cuda::cudaComplex const& b)
   {
      Prdc::Cuda::cudaComplex z;
      sub(z, a, b);
      return absSq<Prdc::Cuda::cudaComplex, Prdc::Cuda::cudaReal>(z);
   }
 
   // Multiplication
 
   template <> inline
   void mul(Prdc::Cuda::cudaComplex& z, Prdc::Cuda::cudaComplex const& a, 
            Prdc::Cuda::cudaComplex const& b)
   {
      z.x = a.x * b.x - a.y * b.y;
      z.y = a.y * b.x + a.x * b.y;
   }
 
   template <> inline
   void mul(Prdc::Cuda::cudaComplex& z, Prdc::Cuda::cudaComplex const& a, 
            Prdc::Cuda::cudaReal const& b)
   {
      z.x = a.x*b;
      z.y = a.y*b;
   }
 
   template <> inline
   void mulEq(Prdc::Cuda::cudaComplex & a, Prdc::Cuda::cudaComplex const& b)
   {
      Prdc::Cuda::cudaReal a0;
      a0   = a.x * b.x - a.y * b.y;
      a.y = a.y * b.x + a.x * b.y;
      a.x = a0;
   }
 
   template <> inline
   void mulEq(Prdc::Cuda::cudaComplex & a, Prdc::Cuda::cudaReal const& b)
   {
      a.x *= b;
      a.y *= b;
   }
 
   template <> inline
   void square(Prdc::Cuda::cudaComplex& z, Prdc::Cuda::cudaComplex const& a)
   {
      z.x = a.x * a.x - a.y * a.y;
      z.y = 2.0 * a.y * a.x;
   }
 
   // Division
 
   template <> inline
   void div(Prdc::Cuda::cudaComplex& z, Prdc::Cuda::cudaComplex const& a, 
            Prdc::Cuda::cudaComplex const& b)
   {
      Prdc::Cuda::cudaReal bSq = b.x * b.x + b.y * b.y;
      z.x = (a.x * b.x + a.y * b.y) / bSq;
      z.y = (a.y * b.x - a.x * b.y) / bSq;
   }
 
   template <> inline
   void div(Prdc::Cuda::cudaComplex& z, Prdc::Cuda::cudaComplex const& a, 
            Prdc::Cuda::cudaReal const& b)
   {
      z.x = a.x / b;
      z.y = a.y / b;
   }
 
   template <> inline
   void divEq(Prdc::Cuda::cudaComplex & a, Prdc::Cuda::cudaComplex const & b)
   {
      Prdc::Cuda::cudaReal bSq = b.x * b.x + b.y * b.y;
      Prdc::Cuda::cudaReal a0 = (a.x * b.x + a.y * b.y)/bSq;
      a.y = (a.y * b.x - a.x * b.y)/bSq;
      a.x = a0;
   }
 
   template <> inline
   void divEq(Prdc::Cuda::cudaComplex & a, Prdc::Cuda::cudaReal const& b)
   {
      a.x /= b;
      a.y /= b;
   }
 
} // Pscf
#endif