Prepare for radial basis.

8 files changed, 128 insertions, 155 deletions

diff --git a/brill_data.h b/brill_data.h
index e928db9..621998d 100644
--- a/brill_data.h
+++ b/brill_data.h
@@ -1,5 +1,5 @@
 /*
- * Copyright 2014 Anton Khirnov <anton@khirnov.net>
+ * Copyright 2014-2015 Anton Khirnov <anton@khirnov.net>
  *
  * This program is free software: you can redistribute it and/or modify
  * it under the terms of the GNU General Public License as published by
@@ -51,9 +51,26 @@ typedef struct BDContext {
      */
     void *priv;
 
-    /*******************************
-     *  options, set by the caller *
-     *******************************/
+    /********************************
+     *  options, set by the caller  *
+     ********************************/
+
+    /**
+     * A callback that will be used to print diagnostic messages.
+     *
+     * Defaults to fprintf(stderr, ...)
+     */
+    void (*log_callback)(const struct BDContext *bd, int level,
+                         const char *fmt, va_list);
+
+    /**
+     * Arbitrary user data, e.g. to be used by the log callback.
+     */
+    void *opaque;
+
+    /********************************
+     *  initial data parameters     *
+     ********************************/
 
     /**
      * The choice of the q function in the exponent.
@@ -72,54 +89,19 @@ typedef struct BDContext {
      */
     unsigned int eppley_n;
 
-    /* the solver parameters */
-
-    /**
-     * The number of basis functions in the ρ direction.
-     * Defaults to 80.
-     */
-    unsigned int nb_coeffs_rho;
-    /**
-     * The number of basis functions in the z direction. 0 means the same as
-     * nb_coeffs_rho. Defaults to 0.
-     * Using values other than 0 or nb_coeffs_rho is unsupported for now.
-     */
-    unsigned int nb_coeffs_z;
-
-    /**
-     * The difference between the number of collocation points and the number
-     * of basis functions in the rho direction. I.e. the number of collocation
-     * points used will be nb_coeffs_rho + overdet_rho.
-     * Defaults to 0;
-     */
-    int overdet_rho;
-    /**
-     * Same as overdet_rho, but for the z direction.
-     */
-    int overdet_z;
-
-    /**
-     * The scaling factor used in the basis functions in the ρ direction.
-     * Defaults to 3.
-     */
-    double basis_scale_factor_rho;
-    /**
-     * Same as basis_scale_factor_rho, but for the z direction
-     */
-    double basis_scale_factor_z;
+    /********************************
+     *        solver options        *
+     ********************************/
 
     /**
-     * A callback that will be used to print diagnostic messages.
-     *
-     * Defaults to fprintf(stderr, ...)
+     * The number of basis functions in each direction.
      */
-    void (*log_callback)(const struct BDContext *bd, int level,
-                         const char *fmt, va_list);
+    unsigned int nb_coeffs[2];
 
     /**
-     * Arbitrary user data, e.g. to be used by the log callback.
+     * The scaling factor used in the basis functions.
      */
-    void *opaque;
+    double basis_scale_factor[2];
 
     /**********
      * output *
diff --git a/brill_data.py b/brill_data.py
index 398a8d8..9bfe25d 100644
--- a/brill_data.py
+++ b/brill_data.py
@@ -39,17 +39,13 @@ class BrillData(object):
 
     class _BDContext(ctypes.Structure):
      _fields_ = [("priv",                   ctypes.c_void_p),
+                 ("log_callback",           ctypes.c_void_p),
+                 ("opaque",                 ctypes.c_void_p),
                  ("q_func_type",            ctypes.c_int),
                  ("amplitude",              ctypes.c_double),
                  ("eppley_n",               ctypes.c_uint),
-                 ("nb_coeffs_rho",          ctypes.c_uint),
-                 ("nb_coeffs_z",            ctypes.c_uint),
-                 ("overdet_rho",            ctypes.c_int),
-                 ("overdet_z",              ctypes.c_int),
-                 ("basis_scale_factor_rho", ctypes.c_double),
-                 ("basis_scale_factor_z",   ctypes.c_double),
-                 ("log_callback",           ctypes.c_void_p),
-                 ("opaque",                 ctypes.c_void_p),
+                 ("nb_coeffs",              ctypes.c_uint * 2),
+                 ("basis_scale_factor",     ctypes.c_double * 2),
                  ("psi_minus1_coeffs",      ctypes.POINTER(ctypes.c_double)),
                  ("stride",                 ctypes.c_long)]
 
@@ -57,13 +53,21 @@ class BrillData(object):
         self._libbd = ctypes.CDLL('libbrilldata.so')
         self._bdctx = ctypes.cast(self._libbd.bd_context_alloc(), ctypes.POINTER(self._BDContext))
         for arg, value in kwargs.iteritems():
-            self._bdctx.contents.__setattr__(arg, value)
+            try:
+                self._bdctx.contents.__setattr__(arg, value)
+            except TypeError as e:
+                # try assigning items of an iterable
+                try:
+                    for i, it in enumerate(value):
+                        self._bdctx.contents.__getattribute__(arg)[i] = it
+                except:
+                    raise e
 
         ret = self._libbd.bd_solve(self._bdctx)
         if ret < 0:
             raise RuntimeError('Error solving the equation')
 
-        self.coeffs = np.copy(np.ctypeslib.as_array(self._bdctx.contents.psi_minus1_coeffs, (self._bdctx.contents.nb_coeffs_z, self._bdctx.contents.nb_coeffs_rho)))
+        self.coeffs = np.copy(np.ctypeslib.as_array(self._bdctx.contents.psi_minus1_coeffs, (self._bdctx.contents.nb_coeffs[1], self._bdctx.contents.nb_coeffs[0])))
 
     def __del__(self):
         if self._bdctx:
diff --git a/eval.c b/eval.c
index deb9579..6c70790 100644
--- a/eval.c
+++ b/eval.c
@@ -35,7 +35,7 @@ int bd_eval_psi(const BDContext *bd,
                 double *psi, unsigned int psi_stride)
 {
     BDPriv *s = bd->priv;
-    const double sfr = bd->basis_scale_factor_rho, sfz = bd->basis_scale_factor_z;
+    const double *sf = bd->basis_scale_factor;
     double (*eval)(double coord, int idx, double sf);
 
     double *basis_val_rho = NULL, *basis_val_z = NULL;
@@ -53,34 +53,34 @@ int bd_eval_psi(const BDContext *bd,
         add = 1.0;
 
     /* precompute the basis values on the grid points */
-    basis_val_rho = malloc(sizeof(*basis_val_rho) * bd->nb_coeffs_rho * nb_coords_rho);
-    basis_val_z   = malloc(sizeof(*basis_val_z)   * bd->nb_coeffs_z   * nb_coords_z);
+    basis_val_rho = malloc(sizeof(*basis_val_rho) * s->nb_coeffs[0] * nb_coords_rho);
+    basis_val_z   = malloc(sizeof(*basis_val_z)   * s->nb_coeffs[1] * nb_coords_z);
     if (!basis_val_rho || !basis_val_z) {
         ret = -ENOMEM;
         goto fail;
     }
 
     switch (diff_order[0]) {
-    case 0: eval = s->basis->eval;       break;
-    case 1: eval = s->basis->eval_diff1; break;
-    case 2: eval = s->basis->eval_diff2; break;
+    case 0: eval = s->basis[0]->eval;       break;
+    case 1: eval = s->basis[0]->eval_diff1; break;
+    case 2: eval = s->basis[0]->eval_diff2; break;
     }
 
     for (int i = 0; i < nb_coords_rho; i++) {
         double rrho = rho[i];
-        for (int j = 0; j < bd->nb_coeffs_rho; j++)
-            basis_val_rho[i * bd->nb_coeffs_rho + j] = eval(rrho, j, sfr);
+        for (int j = 0; j < s->nb_coeffs[0]; j++)
+            basis_val_rho[i * s->nb_coeffs[0] + j] = eval(rrho, j, sf[0]);
     }
 
     switch (diff_order[1]) {
-    case 0: eval = s->basis->eval;       break;
-    case 1: eval = s->basis->eval_diff1; break;
-    case 2: eval = s->basis->eval_diff2; break;
+    case 0: eval = s->basis[1]->eval;       break;
+    case 1: eval = s->basis[1]->eval_diff1; break;
+    case 2: eval = s->basis[1]->eval_diff2; break;
     }
     for (int i = 0; i < nb_coords_z; i++) {
         double zz   = z[i];
-        for (int j = 0; j < bd->nb_coeffs_z; j++)
-            basis_val_z[i * bd->nb_coeffs_z + j] = eval(zz, j, sfz);
+        for (int j = 0; j < s->nb_coeffs[1]; j++)
+            basis_val_z[i * s->nb_coeffs[1] + j] = eval(zz, j, sf[1]);
     }
 
     for (int i = 0; i < nb_coords_z; i++) {
@@ -89,9 +89,9 @@ int bd_eval_psi(const BDContext *bd,
         for (int j = 0; j < nb_coords_rho; j++) {
             double ppsi = add;
 
-            for (int m = 0; m < bd->nb_coeffs_z; m++)
-                for (int n = 0; n < bd->nb_coeffs_rho; n++)
-                    ppsi += s->coeffs[m * bd->nb_coeffs_rho + n] * basis_val_rho[j * bd->nb_coeffs_rho + n] * basis_val_z[i * bd->nb_coeffs_z + m];
+            for (int m = 0; m < s->nb_coeffs[1]; m++)
+                for (int n = 0; n < s->nb_coeffs[0]; n++)
+                    ppsi += s->coeffs[m * s->nb_coeffs[0] + n] * basis_val_rho[j * s->nb_coeffs[0] + n] * basis_val_z[i * s->nb_coeffs[1] + m];
 
             *dst++ = ppsi;
         }
diff --git a/init.c b/init.c
index 759088f..6ea45de 100644
--- a/init.c
+++ b/init.c
@@ -45,20 +45,13 @@ static int brill_init_check_options(BDContext *bd)
         return -EINVAL;
     }
 
-    if (!bd->nb_coeffs_z)
-        bd->nb_coeffs_z = bd->nb_coeffs_rho;
+    s->basis[0] = &bdi_sb_even_basis;
+    s->basis[1] = &bdi_sb_even_basis;
 
-    s->basis = &bdi_sb_even_basis;
-
-    s->nb_coeffs = bd->nb_coeffs_rho * bd->nb_coeffs_z;
-
-    s->nb_colloc_points_rho = bd->nb_coeffs_rho + bd->overdet_rho;
-    s->nb_colloc_points_z   = bd->nb_coeffs_z   + bd->overdet_z;
-
-    s->nb_colloc_points     = s->nb_colloc_points_rho * s->nb_colloc_points_z;
-
-    s->colloc_grid_order_rho = s->nb_colloc_points_rho;
-    s->colloc_grid_order_z   = s->nb_colloc_points_z;
+    s->nb_colloc_points[0] = bd->nb_coeffs[0];
+    s->nb_colloc_points[1] = bd->nb_coeffs[1];
+    s->nb_coeffs[0]        = bd->nb_coeffs[0];
+    s->nb_coeffs[1]        = bd->nb_coeffs[1];
 
     return 0;
 }
@@ -82,7 +75,7 @@ int bd_solve(BDContext *bd)
         return ret;
 
     bd->psi_minus1_coeffs = s->coeffs;
-    bd->stride            = bd->nb_coeffs_rho;
+    bd->stride            = s->nb_coeffs[0];
 
     return 0;
 }
@@ -98,11 +91,11 @@ BDContext *bd_context_alloc(void)
     bd->amplitude   = 1.0;
     bd->eppley_n    = 5;
 
-    bd->nb_coeffs_rho = 80;
-    bd->nb_coeffs_z   = 0;
+    bd->nb_coeffs[0] = 80;
+    bd->nb_coeffs[1] = 80;
 
-    bd->basis_scale_factor_rho = 3.0;
-    bd->basis_scale_factor_z   = 3.0;
+    bd->basis_scale_factor[0] = 3.0;
+    bd->basis_scale_factor[1] = 3.0;
 
     bd->log_callback = bdi_log_default_callback;
 
diff --git a/internal.h b/internal.h
index 87490e1..e0cbf60 100644
--- a/internal.h
+++ b/internal.h
@@ -55,20 +55,18 @@ typedef struct QFunc {
 } QFunc;
 
 typedef struct BDPriv {
-    const BasisSet *basis;
+    const BasisSet *basis[2];
     const QFunc    *q_func;
 
-    int nb_colloc_points;
-    int nb_colloc_points_rho;
-    int nb_colloc_points_z;
-
-    int colloc_grid_order_rho;
-    int colloc_grid_order_z;
-    int nb_coeffs;
+    int nb_colloc_points[2];
+    int nb_coeffs[2];
 
     double *coeffs;
 } BDPriv;
 
+#define NB_COEFFS(s)        (s->nb_coeffs[0]        * s->nb_coeffs[1])
+#define NB_COLLOC_POINTS(s) (s->nb_colloc_points[0] * s->nb_colloc_points[1])
+
 extern const BasisSet bdi_sb_even_basis;
 
 extern const QFunc bdi_q_func_gundlach;
diff --git a/libbrilldata.v b/libbrilldata.v
index c6cfab2..74587fd 100644
--- a/libbrilldata.v
+++ b/libbrilldata.v
@@ -1,4 +1,4 @@
-LIBBRILLDATA_0 {
+LIBBRILLDATA_1 {
     global: bd_*;
     local: *;
 };
diff --git a/solve.c b/solve.c
index ab8fc39..b9b26a2 100644
--- a/solve.c
+++ b/solve.c
@@ -37,80 +37,76 @@ static int brill_construct_matrix(const BDContext *bd, double *mat,
                                   double *rhs)
 {
     BDPriv *s = bd->priv;
-    const double sfr = bd->basis_scale_factor_rho, sfz = bd->basis_scale_factor_z;
+    const double *sf = bd->basis_scale_factor;
 
     double *basis_val, *basis_dval, *basis_d2val;
     int idx_coeff_rho, idx_coeff_z, idx_grid_rho, idx_grid_z;
 
-    double *basis_r[3];
-    double *basis_z[3];
+    double *basis[2][3] = { { NULL } };
+    int ret = 0;
 
     /* precompute the basis values we will need */
-    for (int i = 0; i < ARRAY_ELEMS(basis_r); i++)
-        basis_r[i] = malloc(sizeof(*basis_r[i]) * s->nb_colloc_points_rho * bd->nb_coeffs_rho);
-    for (int i = 0; i < s->nb_colloc_points_rho; i++) {
-        double coord = s->basis->colloc_point(s->colloc_grid_order_rho, i, sfr);
-        for (int j = 0; j < bd->nb_coeffs_rho; j++) {
-            basis_r[0][i * bd->nb_coeffs_rho + j] = s->basis->eval(coord, j, sfr);
-            basis_r[1][i * bd->nb_coeffs_rho + j] = s->basis->eval_diff1(coord, j, sfr);
-            basis_r[2][i * bd->nb_coeffs_rho + j] = s->basis->eval_diff2(coord, j, sfr);
+    for (int i = 0; i < ARRAY_ELEMS(basis); i++) {
+        for (int j = 0; j < ARRAY_ELEMS(basis[i]); j++) {
+            basis[i][j] = malloc(sizeof(*basis[i][j]) * s->nb_colloc_points[i] * s->nb_coeffs[i]);
+            if (!basis[i][j]) {
+                ret = -ENOMEM;
+                goto fail;
+            }
         }
-    }
-
-    for (int i = 0; i < ARRAY_ELEMS(basis_z); i++)
-        basis_z[i] = malloc(sizeof(*basis_z[i]) * s->nb_colloc_points_z * bd->nb_coeffs_z);
-    for (int i = 0; i < s->nb_colloc_points_z; i++) {
-        double coord = s->basis->colloc_point(s->colloc_grid_order_z, i, sfr);
-        for (int j = 0; j < bd->nb_coeffs_z; j++) {
-            basis_z[0][i * bd->nb_coeffs_z + j] = s->basis->eval(coord, j, sfr);
-            basis_z[1][i * bd->nb_coeffs_z + j] = s->basis->eval_diff1(coord, j, sfr);
-            basis_z[2][i * bd->nb_coeffs_z + j] = s->basis->eval_diff2(coord, j, sfr);
+        for (int j = 0; j < s->nb_colloc_points[i]; j++) {
+            double coord = s->basis[i]->colloc_point(s->nb_coeffs[i], j, sf[i]);
+            for (int k = 0; k < s->nb_coeffs[i]; k++) {
+                basis[i][0][j * s->nb_coeffs[i] + k] = s->basis[i]->eval      (coord, k, sf[i]);
+                basis[i][1][j * s->nb_coeffs[i] + k] = s->basis[i]->eval_diff1(coord, k, sf[i]);
+                basis[i][2][j * s->nb_coeffs[i] + k] = s->basis[i]->eval_diff2(coord, k, sf[i]);
+            }
         }
     }
 
-#define   BASIS_RHO (basis_r[0][idx_grid_rho * bd->nb_coeffs_rho + idx_coeff_rho])
-#define  DBASIS_RHO (basis_r[1][idx_grid_rho * bd->nb_coeffs_rho + idx_coeff_rho])
-#define D2BASIS_RHO (basis_r[2][idx_grid_rho * bd->nb_coeffs_rho + idx_coeff_rho])
-#define   BASIS_Z   (basis_z[0][idx_grid_z * bd->nb_coeffs_z + idx_coeff_z])
-#define  DBASIS_Z   (basis_z[1][idx_grid_z * bd->nb_coeffs_z + idx_coeff_z])
-#define D2BASIS_Z   (basis_z[2][idx_grid_z * bd->nb_coeffs_z + idx_coeff_z])
+#define   BASIS_RHO (basis[0][0][idx_grid_rho * s->nb_coeffs[0] + idx_coeff_rho])
+#define  DBASIS_RHO (basis[0][1][idx_grid_rho * s->nb_coeffs[0] + idx_coeff_rho])
+#define D2BASIS_RHO (basis[0][2][idx_grid_rho * s->nb_coeffs[0] + idx_coeff_rho])
+#define   BASIS_Z   (basis[1][0][idx_grid_z   * s->nb_coeffs[1] + idx_coeff_z])
+#define  DBASIS_Z   (basis[1][1][idx_grid_z   * s->nb_coeffs[1] + idx_coeff_z])
+#define D2BASIS_Z   (basis[1][2][idx_grid_z   * s->nb_coeffs[1] + idx_coeff_z])
 
-    for (idx_grid_z = 0; idx_grid_z < s->nb_colloc_points_z; idx_grid_z++) {
-        double z_physical = s->basis->colloc_point(s->colloc_grid_order_z, idx_grid_z, sfz);
+    for (idx_grid_z = 0; idx_grid_z < s->nb_colloc_points[1]; idx_grid_z++) {
+        double z_val = s->basis[1]->colloc_point(s->nb_coeffs[1], idx_grid_z, sf[1]);
 
-        for (idx_grid_rho = 0; idx_grid_rho < s->nb_colloc_points_rho; idx_grid_rho++) {
-            double x_physical = s->basis->colloc_point(s->colloc_grid_order_rho, idx_grid_rho, sfr);
-            double d2q        = s->q_func->d2q_rho(bd, x_physical, z_physical) + s->q_func->d2q_z(bd, x_physical, z_physical);
-            int idx_grid         = idx_grid_z * s->nb_colloc_points_rho + idx_grid_rho;
+        for (idx_grid_rho = 0; idx_grid_rho < s->nb_colloc_points[0]; idx_grid_rho++) {
+            double x_val = s->basis[0]->colloc_point(s->nb_coeffs[0], idx_grid_rho, sf[0]);
+            double d2q   = s->q_func->d2q_rho(bd, x_val, z_val) + s->q_func->d2q_z(bd, x_val, z_val);
+            int idx_grid = idx_grid_z * s->nb_colloc_points[0] + idx_grid_rho;
 
-            for (idx_coeff_z = 0; idx_coeff_z < bd->nb_coeffs_z; idx_coeff_z++)
-                for (idx_coeff_rho = 0; idx_coeff_rho < bd->nb_coeffs_rho; idx_coeff_rho++) {
-                    int idx_coeff = idx_coeff_z * bd->nb_coeffs_rho + idx_coeff_rho;
+            for (idx_coeff_z = 0; idx_coeff_z < s->nb_coeffs[1]; idx_coeff_z++)
+                for (idx_coeff_rho = 0; idx_coeff_rho < s->nb_coeffs[0]; idx_coeff_rho++) {
+                    int idx_coeff = idx_coeff_z * s->nb_coeffs[0] + idx_coeff_rho;
 
                     double val  = D2BASIS_RHO * BASIS_Z + D2BASIS_Z * BASIS_RHO + BASIS_RHO * BASIS_Z * 0.25 * d2q;
-                    if (fabs(x_physical) > EPS)
-                        val += DBASIS_RHO * BASIS_Z / fabs(x_physical);
+                    if (fabs(x_val) > EPS)
+                        val += DBASIS_RHO * BASIS_Z / fabs(x_val);
                     else
                         val += D2BASIS_RHO * BASIS_Z;
 
-                    mat[idx_grid + s->nb_colloc_points * idx_coeff] = val;
+                    mat[idx_grid + NB_COLLOC_POINTS(s) * idx_coeff] = val;
                 }
             rhs[idx_grid] = -0.25 * d2q;
         }
     }
 
-    for (int i = 0; i < ARRAY_ELEMS(basis_r); i++)
-        free(basis_r[i]);
-    for (int i = 0; i < ARRAY_ELEMS(basis_z); i++)
-        free(basis_z[i]);
+fail:
+    for (int i = 0; i < ARRAY_ELEMS(basis); i++)
+        for (int j = 0; j < ARRAY_ELEMS(basis[i]); j++)
+            free(basis[i][j]);
 
-    return 0;
+    return ret;
 }
 
 static int brill_solve_linear(const BDContext *bd, double *mat, double **px, double **prhs)
 {
     const BDPriv *s = bd->priv;
-    const int stride = s->nb_coeffs;
+    const int stride = NB_COEFFS(s);
     int *ipiv;
     double *mat_f;
     double cond, ferr, berr, rpivot;
@@ -152,18 +148,18 @@ static int brill_solve_svd(const BDContext *bd, double *mat,
     double   *x = *px;
     double *rhs = *prhs;
 
-    sv = malloc(sizeof(*sv) * s->nb_coeffs);
+    sv = malloc(sizeof(*sv) * NB_COEFFS(s));
     if (!sv)
         return -ENOMEM;
 
-    LAPACKE_dgelsd(LAPACK_COL_MAJOR, s->nb_colloc_points, s->nb_coeffs, 1,
-                   mat, s->nb_colloc_points, rhs,
-                   MAX(s->nb_coeffs, s->nb_colloc_points),
+    LAPACKE_dgelsd(LAPACK_COL_MAJOR, NB_COLLOC_POINTS(s), NB_COEFFS(s), 1,
+                   mat, NB_COLLOC_POINTS(s), rhs,
+                   MAX(NB_COEFFS(s), NB_COLLOC_POINTS(s)),
                    sv, -1, &rank);
 
     bdi_log(bd, 0, "Least squares SVD solution to a %zdx%zd matrix: "
-            "rank %d, condition number %16.16g\n", s->nb_colloc_points, s->nb_coeffs,
-            rank, sv[s->nb_coeffs - 1] / sv[0]);
+            "rank %d, condition number %16.16g\n", NB_COLLOC_POINTS(s), NB_COEFFS(s),
+            rank, sv[NB_COEFFS(s) - 1] / sv[0]);
 
     free(sv);
 
@@ -189,7 +185,7 @@ int bdi_solve(BDContext *bd)
 {
     BDPriv *s = bd->priv;
 
-    const int vecsize = MAX(s->nb_coeffs, s->nb_colloc_points);
+    const int vecsize = MAX(NB_COEFFS(s), NB_COLLOC_POINTS(s));
 
     double *mat = NULL;
     double *rhs = NULL, *coeffs = NULL;
@@ -197,7 +193,7 @@ int bdi_solve(BDContext *bd)
     int ret = 0;
 
     /* allocate and fill the pseudospectral matrix */
-    mat    = malloc(sizeof(*mat) * s->nb_coeffs * s->nb_colloc_points);
+    mat    = malloc(sizeof(*mat) * NB_COEFFS(s) * NB_COLLOC_POINTS(s));
     coeffs = malloc(sizeof(*coeffs) * vecsize);
     rhs    = malloc(sizeof(*rhs)    * vecsize);
     if (!mat || !coeffs || !rhs) {
diff --git a/test.py b/test.py
index 5af80bc..f6e7b51 100755
--- a/test.py
+++ b/test.py
@@ -40,7 +40,7 @@ for amplitude, err_row in zip(amplitudes, err_coeffs):
     for coeff, err_ref in zip(coeffs, err_row):
         sys.stderr.write('coeffs: %dx%d ' % (coeff, coeff))
 
-        BD = brill_data.BrillData(nb_coeffs_rho = coeff, amplitude = amplitude)
+        BD = brill_data.BrillData(nb_coeffs = [coeff, coeff], amplitude = amplitude)
 
         error = np.max(np.abs(eval_residual_coeffs(BD, x, z)))
         sys.stderr.write('calculated error %g, reference %g\n' % (error, err_ref))
@@ -73,7 +73,7 @@ def eval_residual_fd(bd, x, y, z):
 points = [ 100,  500,  1000 ]
 err_fd = [ 1e-2, 5e-4, 1e-4 ]
 sys.stderr.write('Testing convergence of the residual calculated by finite difference derivatives\n')
-BD = brill_data.BrillData(nb_coeffs_rho = 50, amplitude = 5)
+BD = brill_data.BrillData(nb_coeffs = [50, 50], amplitude = 5)
 for point, err_ref in zip(points, err_fd):
     x    = np.linspace(-1, 1, point)
     y    = np.linspace(-3 * (x[1] - x[0]), 3 * (x[1] - x[0]), 7)
@@ -97,7 +97,7 @@ err_metric = [
     [ 5e-4, 5e-5, 5e-4, 1e-4, 1e-4 ],    # component 22
 ]
 
-BD = brill_data.BrillData(nb_coeffs_rho = 40, amplitude = 3)
+BD = brill_data.BrillData(nb_coeffs = [40, 40], amplitude = 3)
 for c, err_row in zip(xrange(3), err_metric):
     sys.stderr.write('component %d%d\n' % (c, c))