cuda/complex.h

#ifndef PSCF_CUDA_COMPLEX_H
#define PSCF_CUDA_COMPLEX_H
 
/*
* PSCF - Polymer Self-Consistent Field
*
* Copyright 2015 - 2025, The Regents of the University of Minnesota
* Distributed under the terms of the GNU General Public License.
*/
 
#include <pscf/cuda/cudaTypes.h>
#include <pscf/math/arithmetic.h>
#include <complex>
 
namespace Pscf {

 
   /*
   * Types cudaComplex and cudaReal are defined in pscf/cuda/cudaTypes.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

   inline 
   cudaReal real(cudaComplex const& a)
   {  return a.x; }

   inline 
   cudaReal imag(cudaComplex const& a)
   {  return a.y; }
 
   // Absolute magnitude

   inline 
   cudaReal abs(cudaComplex const& a)
   {  return sqrt(a.x * a.x + a.y * a.y); }

   inline 
   cudaReal absSq(cudaComplex const& a)
   {  return (a.x * a.x + a.y * a.y); }
 
   // Complex Conjugation

   inline
   void conj(cudaComplex& z, cudaComplex const& a)
   {
      z.x = a.x;
      z.y = -a.y;
   }

   inline
   void conj(cudaComplex& a)
   {
      a.x = a.x;
      a.y = -a.y;
   }
 
   // Assignment

   inline
   void assign(cudaComplex& z, cudaReal const& a, cudaReal const& b)
   {
      z.x = a;
      z.y = b;
   }

   inline
   void assign(cudaComplex& z, cudaReal const& a)
   {
      z.x = a;
      z.y = 0;
   }

   inline
   void assign(cudaComplex& z, cudaComplex const& a)
   {
      z.x = a.x;
      z.y = a.y;
   }

   inline
   void assign(cudaComplex & z, 
               std::complex<cudaReal> const& a)
   {
      z.x = a.real();
      z.y = a.imag();
   }

   inline
   void assign(std::complex<cudaReal> & z, 
               cudaComplex const& a)
   {  z = std::complex<cudaReal>(a.x, a.y); }
 
   // Addition

   inline
   void add(cudaComplex& z, cudaComplex const& a, 
            cudaComplex const& b)
   {
      z.x = a.x + b.x;
      z.y = a.y + b.y;
   }

   inline
   void add(cudaComplex& z, cudaComplex const& a, 
            cudaReal const& b)
   {
      z.x = a.x + b;
      z.y = a.y;
   }

   inline
   void addEq(cudaComplex& a, cudaComplex const& b)
   {
      a.x += b.x;
      a.y += b.y;
   }

   inline
   void addEq(cudaComplex& a, cudaReal const& b)
   {
      a.x += b;
   }
 
   // Subtraction

   inline
   void sub(cudaComplex& z, cudaComplex const& a, 
            cudaComplex const& b)
   {
      z.x = a.x - b.x;
      z.y = a.y - b.y;
   }

   inline
   void sub(cudaComplex& z, cudaComplex const& a, 
            cudaReal const& b)
   {
      z.x = a.x - b;
      z.y = a.y;
   }

   inline
   void subEq(cudaComplex & a, cudaComplex const& b)
   {
      a.x -= b.x;
      a.y -= b.y;
   }

   inline
   void subEq(cudaComplex & a, cudaReal const& b)
   {
      a.x -= b;
   }

   inline 
   cudaReal absSqDiff(cudaComplex const& a, 
                      cudaComplex const& b)
   {
      cudaComplex z;
      sub(z, a, b);
      return absSq(z);
   }
 
   // Multiplication

   inline
   void mul(cudaComplex& z, cudaComplex const& a, 
            cudaComplex const& b)
   {
      z.x = a.x * b.x - a.y * b.y;
      z.y = a.y * b.x + a.x * b.y;
   }

   inline
   void mul(cudaComplex& z, cudaComplex const& a, 
            cudaReal const& b)
   {
      z.x = a.x*b;
      z.y = a.y*b;
   }

   inline
   void mulEq(cudaComplex & a, cudaComplex const& b)
   {
      cudaReal a0;
      a0   = a.x * b.x - a.y * b.y;
      a.y = a.y * b.x + a.x * b.y;
      a.x = a0;
   }

   inline
   void mulEq(cudaComplex & a, cudaReal const& b)
   {
      a.x *= b;
      a.y *= b;
   }

   inline
   void square(cudaComplex& z, cudaComplex const& a)
   {
      z.x = a.x * a.x - a.y * a.y;
      z.y = 2.0 * a.y * a.x;
   }
 
   // Division

   inline
   void div(cudaComplex& z, cudaComplex const& a, 
            cudaComplex const& b)
   {
      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;
   }

   inline
   void div(cudaComplex& z, cudaComplex const& a, 
            cudaReal const& b)
   {
      z.x = a.x / b;
      z.y = a.y / b;
   }

   inline
   void divEq(cudaComplex & a, cudaComplex const & b)
   {
      cudaReal bSq = b.x * b.x + b.y * b.y;
      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;
   }

   inline
   void divEq(cudaComplex & a, cudaReal const& b)
   {
      a.x /= b;
      a.y /= b;
   }
 
   // Inversion

   inline
   void inverse(cudaComplex& z, 
                cudaComplex const & a)
   {
      cudaReal aSq = a.x * a.x + a.y * a.y;
      z.x =   a.x/aSq;
      z.y = - a.y/aSq;
   }
 
   // Exponentiation and logarithm

   inline
   void assignExp(cudaComplex & z, 
                  cudaComplex const & a)
   {  
      std::complex<cudaReal> arg 
                  = std::complex<cudaReal>(a.x, a.y); 
      std::complex<cudaReal> result = std::exp(arg);
      z.x = result.real();
      z.y = result.imag();
   }

   inline
   void assignLog(cudaComplex & z, 
                  cudaComplex const & a)
   {  
      std::complex<cudaReal> arg 
                 = std::complex<cudaReal>(a.x, a.y); 
      std::complex<cudaReal> result = std::log(arg);
      z.x = result.real();
      z.y = result.imag();
   }
 
   // Pseudo-constructor
 
   /*
   * Pseudo-constructor function for cudaComplex.
   *
   * \ingroup Pscf_Cuda_Complex_Module
   *
   * \param x  real part (in)
   * \param y  imaginary part (in)
   */
   __host__ __device__ static inline   
   cudaComplex makeComplex(cudaReal x, cudaReal y)
   {
      cudaComplex result;
      result.x = x;
      result.y = y;
      return result;
   }
 
} // namespace Pscf
#endif