rz_poly.cc

Go to the documentation of this file.

#include <iostream>
#include "rz_poly.h"

using namespace std;
using namespace magfieldparam;

//_______________________________________________________________________________
rz_poly::rz_poly(int N)
{
   int nz, nr = 0, nv, nt;
   poly_term  v3;
   
   if (N < 2) N = 2;
   data.reserve(N);

   v3.coeff = 1.;
   v3.np[0] = 0;
   v3.np[1] = 1;
   
   data.push_back(poly_vect(1, v3));
   
   for (int m = 2; m <=N; ++m) {
      nz = m;
      nr = 0;
      nv = 0;
      v3.coeff = 1./m;
      v3.np[0] = nr;
      v3.np[1] = nz;

      nt = (m + 2) / 2;
      poly_vect v3x(nt, v3);
      
      while (nz >= 2) {
         nz -= 2;
         nr += 2;
         nv += 1;
         v3x[nv].coeff = -v3x[nv-1].coeff*(nz+1)*(nz+2)/(nr*nr);
         v3x[nv].np[0] = nr;
         v3x[nv].np[1] = nz;
      }
      data.push_back(v3x);
   }
   max_nr = nr+1;
   max_nz = N +1;
   r_pow = new double [max_nr];
   z_pow = new double [max_nz];
   fill(r_pow, r_pow+max_nr, 1.);
   fill(z_pow, z_pow+max_nz, 1.);
   
   n_active = data.size();
   is_off = new bool [n_active];
   fill(is_off, is_off+n_active, false);
}

//_______________________________________________________________________________
rz_poly::rz_poly(const rz_poly& S)
{
   data = S.data;
   max_nr = S.max_nr;
   max_nz = S.max_nz;
   n_active = S.n_active;

   if (max_nr) {
      r_pow = new double [max_nr];
      copy(S.r_pow, S.r_pow+max_nr, r_pow);
   } else r_pow = nullptr;

   if (max_nz) {
      z_pow = new double [max_nz];
      copy(S.z_pow, S.z_pow+max_nz, z_pow);
   } else z_pow = nullptr;

   if (S.is_off) {
      is_off = new bool [data.size()];
      copy(S.is_off, S.is_off+data.size(), is_off);
   } else is_off = nullptr;
}

//_______________________________________________________________________________
rz_poly::~rz_poly()
{
   if (is_off) delete [] is_off;
   if (r_pow)  delete [] r_pow;
   if (z_pow)  delete [] z_pow;
}

//_______________________________________________________________________________
void rz_poly::SetOFF(int npoly)
{
   if ((npoly < 0) || (npoly >= (int)data.size())) return;
   if (is_off[npoly]) return;
   is_off[npoly] = true;
   --n_active;
}

//_______________________________________________________________________________
void rz_poly::SetON(int npoly)
{
   if ((npoly < 0) || (npoly >= (int)data.size())) return;
   if (is_off[npoly]) {
      is_off[npoly] = false;
      ++n_active;
   }
}

//_______________________________________________________________________________
void rz_poly::Print()
{
   if (data.empty()) {
      cout << "The \"rz_poly\" object is NOT initialized!" << endl;
      return;
   }
   
   for (unsigned int ip = 0; ip < data.size(); ++ip) {
      cout << "Polynomial " << ip << " (size=" << data[ip].size() << ", status=";
      if (is_off[ip]) cout << "OFF):" << endl;
      else            cout << "ON):"  << endl;
      for (unsigned int it = 0; it < data[ip].size(); ++it) {
         cout << "\tnr="    << data[ip][it].np[0]
              << "\tnz="    << data[ip][it].np[1]
              << "\tcoeff=" << data[ip][it].coeff
              << endl;
      }
   }
}

//_______________________________________________________________________________
rz_poly rz_poly::Diff(int nvar, bool keep_empty)
{
//Return the derivative of original polynomial by variable nvar. 
//If keep_empty=true, resulting polynomial may contain zero basis terms,
//otherwise the resulting polynomial will be compressed, so that corresponding
//basis terms will be shifted relatively to the original polynomial.
//
   poly_term  v3;
   rz_poly p_out;
   p_out.data.reserve(data.size());
   
   bool *tmp_mask = new bool [data.size()];
   unsigned int  ind_tmp = 0;
   int   tmp_act = 0;
   
   for (unsigned int ip = 0; ip < data.size(); ++ip) {
      poly_vect v3x;
      v3x.reserve(data[ip].size());
      for (unsigned int it = 0; it < data[ip].size(); ++it) {
         v3 = data[ip][it];
         v3.coeff = data[ip][it].coeff * data[ip][it].np[nvar];
         if (v3.coeff != 0) {
            v3.np[nvar] = data[ip][it].np[nvar] - 1;
            v3x.push_back(v3);
         }
      }
      if (!v3x.empty() || keep_empty) {
         v3x.resize(v3x.size());
         p_out.data.push_back(v3x);
         tmp_mask[ind_tmp] = is_off[ip];
         ++ind_tmp;
         if (! is_off[ip]) ++tmp_act;
      }
   }
   
   p_out.data.resize(p_out.data.size());
   p_out.max_nr = max_nr;
   p_out.max_nz = max_nz;

   if (nvar == 0) --p_out.max_nr; else --p_out.max_nz;
   
   p_out.r_pow = new double [p_out.max_nr];
   copy(r_pow, r_pow+p_out.max_nr, p_out.r_pow);
   p_out.z_pow = new double [p_out.max_nz];
   copy(z_pow, z_pow+p_out.max_nz, p_out.z_pow);
   
   p_out.n_active = tmp_act;
   p_out.is_off = new bool [p_out.data.size()];
   copy(tmp_mask, tmp_mask+p_out.data.size(), p_out.is_off);
   
   delete [] tmp_mask;

   return p_out;
}

//_______________________________________________________________________________
rz_poly rz_poly::Int(int nvar)
{
//Return the integral of original polynomial by variable nvar
//
   rz_poly p_out;
   p_out.data = data;
   for (unsigned int ip = 0; ip < data.size(); ++ip) {
      for (unsigned int it = 0; it < data[ip].size(); ++it) {
         p_out.data[ip][it].coeff /= ++p_out.data[ip][it].np[nvar];
      }
   }
   
   p_out.max_nr = max_nr;
   p_out.max_nz = max_nz;

   if (nvar == 0) ++p_out.max_nr; else ++p_out.max_nz;
   
   p_out.r_pow = new double [p_out.max_nr];
   copy(r_pow, r_pow+max_nr, p_out.r_pow);
   p_out.z_pow = new double [p_out.max_nz];
   copy(z_pow, z_pow+max_nz, p_out.z_pow);
   
   if (nvar == 0) p_out.r_pow[max_nr] = p_out.r_pow[max_nr-1] * r_pow[1];
   else           p_out.z_pow[max_nz] = p_out.z_pow[max_nz-1] * z_pow[1];

   p_out.n_active = n_active;
   p_out.is_off = new bool [data.size()];
   copy(is_off, is_off+data.size(), p_out.is_off);
   
   return p_out;
}

//_______________________________________________________________________________
rz_poly& rz_poly::operator*=(double C)
{
//Multiply the polynomial by a constant. Skips terms that are switched off

   for (unsigned int ip = 0; ip < data.size(); ++ip) {
      if (is_off[ip]) continue;
      for (unsigned int it = 0; it < data[ip].size(); ++it) {
         data[ip][it].coeff *= C;
      }
   }
   return *this;
}

//_______________________________________________________________________________
rz_poly& rz_poly::operator*=(double *C)
{
//Multiply the polynomial by an array. Skips terms that are switched off

   for (unsigned int ip = 0; ip < data.size(); ++ip) {
      if (is_off[ip]) continue;
      for (unsigned int it = 0; it < data[ip].size(); ++it) {
         data[ip][it].coeff *= *C;
      }
      ++C;
   }
   return *this;
}

//_______________________________________________________________________________
double rz_poly::GetSVal(double r, double z, const double *C) const
{
//Return value of a polynomial, ignoring terms, that are switched off

   if (r_pow == nullptr) return 0.;
   
   double term, rez = 0.;
   int ip, it;
   
   if (r != r_pow[1]) for (ip = 1; ip < max_nr; ++ip) r_pow[ip] = r*r_pow[ip-1];
   if (z != z_pow[1]) for (ip = 1; ip < max_nz; ++ip) z_pow[ip] = z*z_pow[ip-1];

   for (ip = 0; ip < (int)data.size(); ++ip) {
      if (is_off[ip]) continue;
      term = 0.;
      for (it = 0; it < (int)data[ip].size(); ++it) {
         term += data[ip][it].coeff
               * r_pow[data[ip][it].np[0]]
               * z_pow[data[ip][it].np[1]];
      }
      rez += *C * term;
      ++C;
   }
   
   return rez;
}

//_______________________________________________________________________________
double *rz_poly::GetVVal(double r, double z, double *rez_out)
{
//return an array of doubleype with values of the basis functions in the point
//(r,z). In a case if rez_out != 0, the rez_out array must be long enough to fit
//in n_active elements. In a case if rez_out == 0, a new array of n_active length
//is created; it is in user's responsibility to free the memory after all;
//
   if (r_pow == nullptr) return nullptr;
   
   double term;
   int ip, it;

   double *rez;
   if (rez_out) rez = rez_out;
   else rez = new double [n_active];

   double *pnt = rez;
   
   if (r != r_pow[1]) for (ip = 1; ip < max_nr; ++ip) r_pow[ip] = r*r_pow[ip-1];
   if (z != z_pow[1]) for (ip = 1; ip < max_nz; ++ip) z_pow[ip] = z*z_pow[ip-1];

   for (ip = 0; ip < (int)data.size(); ++ip) {
      if (is_off[ip]) continue;
      term = 0.;
      for (it = 0; it < (int)data[ip].size(); ++it) {
         term += data[ip][it].coeff
               * r_pow[data[ip][it].np[0]]
               * z_pow[data[ip][it].np[1]];
      }
      *pnt = term;
      ++pnt;
   }
   
   return rez;
}

//_______________________________________________________________________________
double *rz_poly::Expand(double *C)
{
//Take C[n_active] - reduced vector of coefficients and return pointer to 
//expanded vector of coefficients of the data.size() length. It is in user's
//responsibility to free the memory after all;
//
   double *rez = new double [data.size()];
   for (int ip = 0; ip < (int)data.size(); ++ip) {
      if (is_off[ip]) rez[ip] = 0.;
      else {
         rez[ip] = *C;
         ++C;
      }
   }
   return rez;
}