newton.c

   1 // Solve quadratic equations involving continued fractions.
   2 // A slight misnomer: Gosper describes how to use Newton's method, but
   3 // I use binary search instead, and assume unique roots in the right ranges.
   4 //
   5 // For technical reasons (see Gosper), we write quadratics as
   6 //
   7 //     a0 xy + a1 x + a2 y + a3
   8 // y = ------------------------
   9 //     a4 xy - a0 x + a5 y - a2
  10 //
  11 // where y is the variable and x is some fixed continued fraction constant.
  12
  13 #include <stdio.h>
  14 #include <stdlib.h>
  15 #include <gmp.h>
  16 #include "cf.h"
  17
  18 struct newton_data_s {
  19   cf_t x;
  20   mpz_t a[6];
  21   mpz_t lower;
  22 };
  23 typedef struct newton_data_s newton_data_t[1];
  24 typedef struct newton_data_s *newton_data_ptr;
  25
  26 // In the 3D table, we have two sets of equations. Left is a0, b0, c0,
  27 // Right is a1, b1, c1, and these are the terms involving x:
  28 //
  29 //   a0 y + b0     a1 y + b1
  30 //   ---------     ---------
  31 //   c0 y - a0     c1 y - a1
  32 //
  33 // I work left to right, and top to bottom, unlike Gosper.
  34 //
  35 // This notation is different to that in cf_bihom.c, and similar to cf_mobius.
  36
  37 struct abc_s {
  38   mpz_t a0, a1;
  39   mpz_t b0, b1;
  40   mpz_t c0, c1;
  41 };
  42 typedef struct abc_s abc_t[1];
  43
  44 void abc_init(abc_t p) {
  45   mpz_init(p->a0); mpz_init(p->a1);
  46   mpz_init(p->b0); mpz_init(p->b1);
  47   mpz_init(p->c0); mpz_init(p->c1);
  48 }
  49
  50 void abc_clear(abc_t p) {
  51   mpz_clear(p->a0); mpz_clear(p->a1);
  52   mpz_clear(p->b0); mpz_clear(p->b1);
  53   mpz_clear(p->c0); mpz_clear(p->c1);
  54 }
  55
  56 void abc_set_coeff(abc_t p, mpz_t a[6]) {
  57   mpz_set(p->a0, a[2]); mpz_set(p->b0, a[3]); mpz_set(p->c0, a[5]);
  58   mpz_set(p->a1, a[0]); mpz_set(p->b1, a[1]); mpz_set(p->c1, a[4]);
  59 }
  60
  61 void abc_print(abc_t p) {
  62   gmp_printf("%Zd y + %Zd    %Zd y + %Zd\n", p->a0, p->b0, p->a1, p->b1);
  63   gmp_printf("%Zd y - %Zd    %Zd y - %Zd\n", p->c0, p->a0, p->c1, p->a1);
  64   gmp_printf("(%Zd+sqrt(%Zd^2+%Zd*%Zd))/%Zd\n", p->a0, p->a0, p->b0, p->c0, p->c0);
  65   gmp_printf("%Zd*z^2-2*%Zd*z-%Zd\n", p->c0, p->a0, p->b0);
  66 }
  67
  68 // Finds smallest root greater than given lower bound.
  69 // Assumes there exists exactly one root with this property.
  70 // (Impossible to solve equations if roots have same integer part at
  71 // the moment. Don't do that.)
  72 // TODO: Rewrite so you choose if you want smaller or greater root.
  73 // or so it returns both solutions in an array.
  74 static void *newton(cf_t cf) {
  75   newton_data_ptr nd = cf_data(cf);
  76   abc_t p;
  77   abc_init(p);
  78   abc_set_coeff(p, nd->a);
  79   cf_t x = nd->x;
  80   mpz_t z, z0, z1, one, pow2, t0, t1, t2;
  81   mpz_init(z); mpz_init(z0); mpz_init(z1); mpz_init(pow2);
  82   mpz_init(t0); mpz_init(t1); mpz_init(t2);
  83   mpz_init(one);
  84   mpz_set_ui(one, 1);
  85
  86   void move_right() {
  87     cf_get(z, x);
  88     mpz_mul(t0, z, p->b1);
  89     mpz_add(t0, t0, p->b0);
  90     mpz_set(p->b0, p->b1);
  91     mpz_set(p->b1, t0);
  92
  93     mpz_mul(t0, z, p->a1);  mpz_mul(t1, z, p->c1);
  94     mpz_add(t0, t0, p->a0); mpz_add(t1, t1, p->c0);
  95     mpz_set(p->a0, p->a1);  mpz_set(p->c0, p->c1);
  96     mpz_set(p->a1, t0);     mpz_set(p->c1, t1);
  97   }
  98
  99   void move_down() {
 100     // Recurrence relation with z, then
 101     // Subtract and invert z.
 102     mpz_mul(t0, z, p->a0);
 103     mpz_add(t0, t0, p->b0);
 104     mpz_mul(t1, z, p->c0);
 105     mpz_sub(t1, t1, p->a0);
 106     mpz_mul(p->a0, z, t1);
 107     mpz_set(p->b0, p->c0);
 108     mpz_sub(p->c0, t0, p->a0);
 109     mpz_set(p->a0, t1);
 110
 111     mpz_mul(t0, z, p->a1);
 112     mpz_add(t0, t0, p->b1);
 113     mpz_mul(t1, z, p->c1);
 114     mpz_sub(t1, t1, p->a1);
 115     mpz_mul(p->a1, z, t1);
 116     mpz_set(p->b1, p->c1);
 117     mpz_sub(p->c1, t0, p->a1);
 118     mpz_set(p->a1, t1);
 119   }
 120
 121   int sign_quad() {
 122     // Returns sign of c0 z^2 - 2 a0 z - b0
 123     mpz_mul_si(t0, p->a0, -2);
 124     mpz_addmul(t0, p->c0, z);
 125     mpz_mul(t0, t0, z);
 126     mpz_sub(t0, t0, p->b0);
 127     return mpz_sgn(t0);
 128   }
 129
 130   int sign_quad1() {
 131     // Returns sign of c1 z^2 - 2 a1 z - b1
 132     mpz_mul_si(t0, p->a1, -2);
 133     mpz_addmul(t0, p->c1, z);
 134     mpz_mul(t0, t0, z);
 135     mpz_sub(t0, t0, p->b1);
 136     return mpz_sgn(t0);
 137   }
 138
 139   move_right();  // Get rid of pathological cases.
 140   move_right();
 141   move_right();
 142
 143   // Get integer part, starting search from given lower bound.
 144   void binary_search(mpz_ptr lower) {
 145     while (!mpz_sgn(p->c0)) move_right();
 146     for (;;) {
 147       mpz_set(z0, lower);
 148       mpz_set(z, lower);
 149       int sign = sign_quad();
 150       mpz_set_ui(pow2, 1);
 151       for (;;) {
 152         mpz_add(z, z0, pow2);
 153         if (sign_quad() != sign) break;
 154         mpz_mul_2exp(pow2, pow2, 1);
 155       }
 156       mpz_set(z1, z);
 157
 158       for (;;) {
 159         mpz_add(z, z0, z1);
 160         mpz_div_2exp(z, z, 1);
 161         if (!mpz_cmp(z, z0)) break;
 162         if (sign_quad() == sign) {
 163           mpz_set(z0, z);
 164         } else {
 165           mpz_set(z1, z);
 166         }
 167       }
 168       sign = sign_quad1();
 169       mpz_set(z, z1);
 170       if (sign_quad1() != sign) break;
 171       move_right();
 172     }
 173   }
 174
 175   binary_search(nd->lower);
 176   cf_put(cf, z0);
 177   mpz_set(z, z0);
 178   move_down();
 179
 180   while(cf_wait(cf)) {
 181     binary_search(one);
 182     cf_put(cf, z0);
 183     mpz_set(z, z0);
 184     move_down();
 185   }
 186   abc_clear(p);
 187   mpz_clear(z); mpz_clear(z0); mpz_clear(z1); mpz_clear(pow2);
 188   mpz_clear(t0); mpz_clear(t1); mpz_clear(t2);
 189   mpz_clear(one);
 190   for (int i = 0; i < 6; i++) mpz_clear(nd->a[i]);
 191   mpz_clear(nd->lower);
 192   free(nd);
 193   return NULL;
 194 }
 195
 196 cf_t cf_new_newton(cf_t x, mpz_t a[6], mpz_t lower) {
 197   newton_data_ptr p = malloc(sizeof(*p));
 198   p->x = x;
 199   for (int i = 0; i < 6; i++) {
 200     mpz_init(p->a[i]);
 201     mpz_set(p->a[i], a[i]);
 202   }
 203   mpz_init(p->lower);
 204   mpz_set(p->lower, lower);
 205   return cf_new(newton, p);
 206 }
 207
 208 cf_t cf_new_sqrt(cf_t x) {
 209   newton_data_ptr p = malloc(sizeof(*p));
 210   p->x = x;
 211   mpz_init(p->lower);
 212   for (int i = 0; i < 6; i++) mpz_init(p->a[i]);
 213   // Solve y = x/y.
 214   mpz_set_si(p->a[1], 1);
 215   mpz_set_si(p->a[5], 1);
 216   return cf_new(newton, p);
 217 }
 218
 219 struct newton_int_data_s {
 220   mpz_t coeff[3];
 221   mpz_t lower;
 222 };
 223 typedef struct newton_int_data_s *newton_int_data_ptr;
 224 typedef struct newton_int_data_s newton_int_data_t[1];
 225
 226 // Finds smallest root greater than given lower bound of quadratic equation
 227 // with integer coefficients. If the coefficient of x is odd, we double
 228 // all the coefficients in our representation so we can find integers a, b, c
 229 // such that the quadratic may be written: c y^2 - 2 a y - b = 0.
 230 //
 231 // Assumes there exists exactly one root above given bound.
 232 // (Impossible to solve equations if roots have same integer part at
 233 // the moment. Don't do that.)
 234 // TODO: Rewrite so you choose if you want smaller or greater root.
 235 // or so it returns both solutions in an array.
 236 static void *newton_integer(cf_t cf) {
 237   newton_int_data_ptr p = cf_data(cf);
 238   mpz_t a, b, c;
 239   mpz_init(a); mpz_init(b); mpz_init(c);
 240   mpz_set(a, p->coeff[0]);
 241   mpz_set(b, p->coeff[1]);
 242   mpz_set(c, p->coeff[2]);
 243
 244   mpz_t z, z0, z1, one, pow2, t0, t1, t2;
 245   mpz_init(z); mpz_init(z0); mpz_init(z1); mpz_init(pow2);
 246   mpz_init(t0); mpz_init(t1); mpz_init(t2);
 247   mpz_init(one);
 248   mpz_set_ui(one, 1);
 249
 250   void move_down() {
 251     // Recurrence relation with z, then
 252     // Subtract and invert z.
 253     mpz_mul(t0, z, a);
 254     mpz_add(t0, t0, b);
 255     mpz_mul(t1, z, c);
 256     mpz_sub(t1, t1, a);
 257     mpz_mul(a, z, t1);
 258     mpz_set(b, c);
 259     mpz_sub(c, t0, a);
 260     mpz_set(a, t1);
 261   }
 262
 263   int sign_quad() {
 264     // Returns sign of c z^2 - 2 a z - b
 265     mpz_mul_si(t0, a, -2);
 266     mpz_addmul(t0, c, z);
 267     mpz_mul(t0, t0, z);
 268     mpz_sub(t0, t0, b);
 269     return mpz_sgn(t0);
 270   }
 271
 272   // Get integer part, starting search from given lower bound.
 273   void binary_search(mpz_ptr lower) {
 274     mpz_set(z0, lower);
 275     mpz_set(z, lower);
 276     int sign = sign_quad();
 277     mpz_set_ui(pow2, 1);
 278     for (;;) {
 279       mpz_add(z, z0, pow2);
 280       if (sign_quad() != sign) break;
 281       mpz_mul_2exp(pow2, pow2, 1);
 282     }
 283     mpz_set(z1, z);
 284
 285     for (;;) {
 286       mpz_add(z, z0, z1);
 287       mpz_div_2exp(z, z, 1);
 288       if (!mpz_cmp(z, z0)) break;
 289       if (sign_quad() == sign) {
 290         mpz_set(z0, z);
 291       } else {
 292         mpz_set(z1, z);
 293       }
 294     }
 295     cf_put(cf, z0);
 296     mpz_set(z, z0);
 297     move_down();
 298   }
 299
 300   binary_search(p->lower);
 301
 302   while(cf_wait(cf)) {
 303     binary_search(one);
 304   }
 305   mpz_clear(z); mpz_clear(z0); mpz_clear(z1); mpz_clear(pow2);
 306   mpz_clear(t0); mpz_clear(t1); mpz_clear(t2);
 307   mpz_clear(one);
 308   for (int i = 0; i < 3; i++) mpz_clear(p->coeff[i]);
 309   mpz_clear(p->lower);
 310   free(p);
 311   return NULL;
 312 }
 313
 314 cf_t cf_new_sqrt_pq(mpz_t zp, mpz_t zq) {
 315   newton_int_data_ptr p = malloc(sizeof(*p));
 316   mpz_init(p->lower);
 317   for (int i = 0; i < 3; i++) mpz_init(p->coeff[i]);
 318   mpz_set(p->coeff[1], zp);
 319   mpz_set(p->coeff[2], zq);
 320   return cf_new(newton_integer, p);
 321 }
 322
 323 cf_t cf_new_sqrt_int(int a, int b) {
 324   mpz_t p, q;
 325   mpz_init(p); mpz_init(q);
 326   mpz_set_si(p, a);
 327   mpz_set_si(q, b);
 328   cf_t res = cf_new_sqrt_pq(p, q);
 329   mpz_clear(p); mpz_clear(q);
 330   return res;
 331 }