ParametrizedEngine/src/rz_poly.cc

0001 #include <iostream>
0002 #include "rz_poly.h"
0003
0004 using namespace std;
0005 using namespace magfieldparam;
0006
0007 //_______________________________________________________________________________
0008 rz_poly::rz_poly(int N) {
0009   int nz, nr = 0, nv, nt;
0010   poly_term v3;
0011
0012   if (N < 2)
0013     N = 2;
0014   data.reserve(N);
0015
0016   v3.coeff = 1.;
0017   v3.np[0] = 0;
0018   v3.np[1] = 1;
0019
0020   data.push_back(poly_vect(1, v3));
0021
0022   for (int m = 2; m <= N; ++m) {
0023     nz = m;
0024     nr = 0;
0025     nv = 0;
0026     v3.coeff = 1. / m;
0027     v3.np[0] = nr;
0028     v3.np[1] = nz;
0029
0030     nt = (m + 2) / 2;
0031     poly_vect v3x(nt, v3);
0032
0033     while (nz >= 2) {
0034       nz -= 2;
0035       nr += 2;
0036       nv += 1;
0037       v3x[nv].coeff = -v3x[nv - 1].coeff * (nz + 1) * (nz + 2) / (nr * nr);
0038       v3x[nv].np[0] = nr;
0039       v3x[nv].np[1] = nz;
0040     }
0041     data.push_back(v3x);
0042   }
0043   max_nr = nr + 1;
0044   max_nz = N + 1;
0045   r_pow = new double[max_nr];
0046   z_pow = new double[max_nz];
0047   fill(r_pow, r_pow + max_nr, 1.);
0048   fill(z_pow, z_pow + max_nz, 1.);
0049
0050   n_active = data.size();
0051   is_off = new bool[n_active];
0052   fill(is_off, is_off + n_active, false);
0053 }
0054
0055 //_______________________________________________________________________________
0056 rz_poly::rz_poly(const rz_poly &S) {
0057   data = S.data;
0058   max_nr = S.max_nr;
0059   max_nz = S.max_nz;
0060   n_active = S.n_active;
0061
0062   if (max_nr) {
0063     r_pow = new double[max_nr];
0064     copy(S.r_pow, S.r_pow + max_nr, r_pow);
0065   } else
0066     r_pow = nullptr;
0067
0068   if (max_nz) {
0069     z_pow = new double[max_nz];
0070     copy(S.z_pow, S.z_pow + max_nz, z_pow);
0071   } else
0072     z_pow = nullptr;
0073
0074   if (S.is_off) {
0075     is_off = new bool[data.size()];
0076     copy(S.is_off, S.is_off + data.size(), is_off);
0077   } else
0078     is_off = nullptr;
0079 }
0080
0081 //_______________________________________________________________________________
0082 rz_poly::~rz_poly() {
0083   if (is_off)
0084     delete[] is_off;
0085   if (r_pow)
0086     delete[] r_pow;
0087   if (z_pow)
0088     delete[] z_pow;
0089 }
0090
0091 //_______________________________________________________________________________
0092 void rz_poly::SetOFF(int npoly) {
0093   if ((npoly < 0) || (npoly >= (int)data.size()))
0094     return;
0095   if (is_off[npoly])
0096     return;
0097   is_off[npoly] = true;
0098   --n_active;
0099 }
0100
0101 //_______________________________________________________________________________
0102 void rz_poly::SetON(int npoly) {
0103   if ((npoly < 0) || (npoly >= (int)data.size()))
0104     return;
0105   if (is_off[npoly]) {
0106     is_off[npoly] = false;
0107     ++n_active;
0108   }
0109 }
0110
0111 //_______________________________________________________________________________
0112 void rz_poly::Print() {
0113   if (data.empty()) {
0114     cout << "The \"rz_poly\" object is NOT initialized!" << endl;
0115     return;
0116   }
0117
0118   for (unsigned int ip = 0; ip < data.size(); ++ip) {
0119     cout << "Polynomial " << ip << " (size=" << data[ip].size() << ", status=";
0120     if (is_off[ip])
0121       cout << "OFF):" << endl;
0122     else
0123       cout << "ON):" << endl;
0124     for (unsigned int it = 0; it < data[ip].size(); ++it) {
0125       cout << "\tnr=" << data[ip][it].np[0] << "\tnz=" << data[ip][it].np[1] << "\tcoeff=" << data[ip][it].coeff
0126            << endl;
0127     }
0128   }
0129 }
0130
0131 //_______________________________________________________________________________
0132 rz_poly rz_poly::Diff(int nvar, bool keep_empty) {
0133   //Return the derivative of original polynomial by variable nvar.
0134   //If keep_empty=true, resulting polynomial may contain zero basis terms,
0135   //otherwise the resulting polynomial will be compressed, so that corresponding
0136   //basis terms will be shifted relatively to the original polynomial.
0137   //
0138   poly_term v3;
0139   rz_poly p_out;
0140   p_out.data.reserve(data.size());
0141
0142   bool *tmp_mask = new bool[data.size()];
0143   unsigned int ind_tmp = 0;
0144   int tmp_act = 0;
0145
0146   for (unsigned int ip = 0; ip < data.size(); ++ip) {
0147     poly_vect v3x;
0148     v3x.reserve(data[ip].size());
0149     for (unsigned int it = 0; it < data[ip].size(); ++it) {
0150       v3 = data[ip][it];
0151       v3.coeff = data[ip][it].coeff * data[ip][it].np[nvar];
0152       if (v3.coeff != 0) {
0153         v3.np[nvar] = data[ip][it].np[nvar] - 1;
0154         v3x.push_back(v3);
0155       }
0156     }
0157     if (!v3x.empty() || keep_empty) {
0158       v3x.resize(v3x.size());
0159       p_out.data.push_back(v3x);
0160       tmp_mask[ind_tmp] = is_off[ip];
0161       ++ind_tmp;
0162       if (!is_off[ip])
0163         ++tmp_act;
0164     }
0165   }
0166
0167   p_out.data.resize(p_out.data.size());
0168   p_out.max_nr = max_nr;
0169   p_out.max_nz = max_nz;
0170
0171   if (nvar == 0)
0172     --p_out.max_nr;
0173   else
0174     --p_out.max_nz;
0175
0176   p_out.r_pow = new double[p_out.max_nr];
0177   copy(r_pow, r_pow + p_out.max_nr, p_out.r_pow);
0178   p_out.z_pow = new double[p_out.max_nz];
0179   copy(z_pow, z_pow + p_out.max_nz, p_out.z_pow);
0180
0181   p_out.n_active = tmp_act;
0182   p_out.is_off = new bool[p_out.data.size()];
0183   copy(tmp_mask, tmp_mask + p_out.data.size(), p_out.is_off);
0184
0185   delete[] tmp_mask;
0186
0187   return p_out;
0188 }
0189
0190 //_______________________________________________________________________________
0191 rz_poly rz_poly::Int(int nvar) {
0192   //Return the integral of original polynomial by variable nvar
0193   //
0194   rz_poly p_out;
0195   p_out.data = data;
0196   for (unsigned int ip = 0; ip < data.size(); ++ip) {
0197     for (unsigned int it = 0; it < data[ip].size(); ++it) {
0198       p_out.data[ip][it].coeff /= ++p_out.data[ip][it].np[nvar];
0199     }
0200   }
0201
0202   p_out.max_nr = max_nr;
0203   p_out.max_nz = max_nz;
0204
0205   if (nvar == 0)
0206     ++p_out.max_nr;
0207   else
0208     ++p_out.max_nz;
0209
0210   p_out.r_pow = new double[p_out.max_nr];
0211   copy(r_pow, r_pow + max_nr, p_out.r_pow);
0212   p_out.z_pow = new double[p_out.max_nz];
0213   copy(z_pow, z_pow + max_nz, p_out.z_pow);
0214
0215   if (nvar == 0)
0216     p_out.r_pow[max_nr] = p_out.r_pow[max_nr - 1] * r_pow[1];
0217   else
0218     p_out.z_pow[max_nz] = p_out.z_pow[max_nz - 1] * z_pow[1];
0219
0220   p_out.n_active = n_active;
0221   p_out.is_off = new bool[data.size()];
0222   copy(is_off, is_off + data.size(), p_out.is_off);
0223
0224   return p_out;
0225 }
0226
0227 //_______________________________________________________________________________
0228 rz_poly &rz_poly::operator*=(double C) {
0229   //Multiply the polynomial by a constant. Skips terms that are switched off
0230
0231   for (unsigned int ip = 0; ip < data.size(); ++ip) {
0232     if (is_off[ip])
0233       continue;
0234     for (unsigned int it = 0; it < data[ip].size(); ++it) {
0235       data[ip][it].coeff *= C;
0236     }
0237   }
0238   return *this;
0239 }
0240
0241 //_______________________________________________________________________________
0242 rz_poly &rz_poly::operator*=(double *C) {
0243   //Multiply the polynomial by an array. Skips terms that are switched off
0244
0245   for (unsigned int ip = 0; ip < data.size(); ++ip) {
0246     if (is_off[ip])
0247       continue;
0248     for (unsigned int it = 0; it < data[ip].size(); ++it) {
0249       data[ip][it].coeff *= *C;
0250     }
0251     ++C;
0252   }
0253   return *this;
0254 }
0255
0256 //_______________________________________________________________________________
0257 double rz_poly::GetSVal(double r, double z, const double *C) const {
0258   //Return value of a polynomial, ignoring terms, that are switched off
0259
0260   if (r_pow == nullptr)
0261     return 0.;
0262
0263   double term, rez = 0.;
0264   int ip, it;
0265
0266   if (r != r_pow[1])
0267     for (ip = 1; ip < max_nr; ++ip)
0268       r_pow[ip] = r * r_pow[ip - 1];
0269   if (z != z_pow[1])
0270     for (ip = 1; ip < max_nz; ++ip)
0271       z_pow[ip] = z * z_pow[ip - 1];
0272
0273   for (ip = 0; ip < (int)data.size(); ++ip) {
0274     if (is_off[ip])
0275       continue;
0276     term = 0.;
0277     for (it = 0; it < (int)data[ip].size(); ++it) {
0278       term += data[ip][it].coeff * r_pow[data[ip][it].np[0]] * z_pow[data[ip][it].np[1]];
0279     }
0280     rez += *C * term;
0281     ++C;
0282   }
0283
0284   return rez;
0285 }
0286
0287 //_______________________________________________________________________________
0288 double *rz_poly::GetVVal(double r, double z, double *rez_out) {
0289   //return an array of doubleype with values of the basis functions in the point
0290   //(r,z). In a case if rez_out != 0, the rez_out array must be long enough to fit
0291   //in n_active elements. In a case if rez_out == 0, a new array of n_active length
0292   //is created; it is in user's responsibility to free the memory after all;
0293   //
0294   if (r_pow == nullptr)
0295     return nullptr;
0296
0297   double term;
0298   int ip, it;
0299
0300   double *rez;
0301   if (rez_out)
0302     rez = rez_out;
0303   else
0304     rez = new double[n_active];
0305
0306   double *pnt = rez;
0307
0308   if (r != r_pow[1])
0309     for (ip = 1; ip < max_nr; ++ip)
0310       r_pow[ip] = r * r_pow[ip - 1];
0311   if (z != z_pow[1])
0312     for (ip = 1; ip < max_nz; ++ip)
0313       z_pow[ip] = z * z_pow[ip - 1];
0314
0315   for (ip = 0; ip < (int)data.size(); ++ip) {
0316     if (is_off[ip])
0317       continue;
0318     term = 0.;
0319     for (it = 0; it < (int)data[ip].size(); ++it) {
0320       term += data[ip][it].coeff * r_pow[data[ip][it].np[0]] * z_pow[data[ip][it].np[1]];
0321     }
0322     *pnt = term;
0323     ++pnt;
0324   }
0325
0326   return rez;
0327 }
0328
0329 //_______________________________________________________________________________
0330 double *rz_poly::Expand(double *C) {
0331   //Take C[n_active] - reduced vector of coefficients and return pointer to
0332   //expanded vector of coefficients of the data.size() length. It is in user's
0333   //responsibility to free the memory after all;
0334   //
0335   double *rez = new double[data.size()];
0336   for (int ip = 0; ip < (int)data.size(); ++ip) {
0337     if (is_off[ip])
0338       rez[ip] = 0.;
0339     else {
0340       rez[ip] = *C;
0341       ++C;
0342     }
0343   }
0344   return rez;
0345 }