Index: /trunk/Ohana/src/opihi/cmd.astro/fitpm_irls.c
===================================================================
--- /trunk/Ohana/src/opihi/cmd.astro/fitpm_irls.c	(revision 38650)
+++ /trunk/Ohana/src/opihi/cmd.astro/fitpm_irls.c	(revision 38650)
@@ -0,0 +1,492 @@
+# include "astro.h"
+# define J2000 51544.5       /* Modified Julian date at standard epoch J2000 */
+
+# define ESCAPE(MSG,...) {			\
+    gprint (GP_ERR, MSG, __VA_ARGS__);		\
+    return FALSE; }
+
+typedef struct {
+  double Ro, dRo;
+  double Do, dDo;
+
+  double uR, duR;
+  double uD, duD;
+  
+  double chisq;
+  int Nfit;
+} PMFit_IRLS;
+
+int FitPMonly_IRLS (PMFit_IRLS *fit, double *X, double *dX, double *Y, double *dY, double *T, int Npts, int VERBOSE);
+int IRLS_converged (PMFit_IRLS *fit);
+
+int fitpm_irls (int argc, char **argv) {
+  
+  int i, N;
+
+  Vector *rvec, *dvec, *tvec, *dRvec, *dDvec;
+
+  Vector *mvec = NULL; // mask vector
+  if ((N = get_argument (argc, argv, "-mask"))) {
+    remove_argument (N, &argc, argv);
+    if ((mvec = SelectVector (argv[N], ANYVECTOR, TRUE)) == NULL) return (FALSE);
+    remove_argument (N, &argc, argv);
+    CastVector (mvec, OPIHI_INT);
+  }
+
+  int VERBOSE = FALSE;
+  if ((N = get_argument (argc, argv, "-v"))) {
+    remove_argument (N, &argc, argv);
+    VERBOSE = TRUE;
+  }
+  if ((N = get_argument (argc, argv, "-vv"))) {
+    remove_argument (N, &argc, argv);
+    VERBOSE = 2;
+  }
+
+  if (argc != 6) {
+    gprint (GP_ERR, "USAGE: fitplx_irls (ra) (dR) (dec) (dD) (mjd) [-mask mask]\n");
+    // what about the errors?
+    return (FALSE);
+  }
+
+  /* select input / output buffers */
+  // if ((buf = SelectBuffer (argv[1], OLDBUFFER, TRUE)) == NULL) return (FALSE);
+  // Nx = buf[0].header.Naxis[0];
+  // Ny = buf[0].header.Naxis[1];
+  
+  if ((rvec = SelectVector (argv[1], ANYVECTOR, TRUE)) == NULL) ESCAPE ("missing vector %s\n", argv[1]);
+  if ((dvec = SelectVector (argv[3], ANYVECTOR, TRUE)) == NULL) ESCAPE ("missing vector %s\n", argv[3]);
+  if ((tvec = SelectVector (argv[5], ANYVECTOR, TRUE)) == NULL) ESCAPE ("missing vector %s\n", argv[5]);
+
+  if ((dRvec = SelectVector (argv[2], ANYVECTOR, TRUE)) == NULL) ESCAPE ("missing vector %s\n", argv[2]);
+  if ((dDvec = SelectVector (argv[4], ANYVECTOR, TRUE)) == NULL) ESCAPE ("missing vector %s\n", argv[4]);
+
+  double *R = rvec->elements.Flt;
+  double *D = dvec->elements.Flt;
+  double *T = tvec->elements.Flt;
+  
+  double *dR = dRvec->elements.Flt;
+  double *dD = dDvec->elements.Flt;
+
+  int *mask = NULL;
+  if (mvec) {
+    mask = mvec->elements.Int;
+  }
+
+  N = tvec->Nelements; // XXX check other lengths
+
+  // find mean values to remove
+  double Npts = 0;
+  double Tmean = 0;
+  double Rmean = 0;
+  double Dmean = 0;
+  double Tmin = +1000000;
+  double Tmax = -1000000;
+  for (i = 0; i < N; i++) {
+    if (mask && !mask[i]) continue;
+    Rmean += R[i];
+    Dmean += D[i];
+    Tmean += T[i];
+    Tmin = MIN(Tmin, T[i]);
+    Tmax = MAX(Tmax, T[i]);
+    Npts += 1.0;
+  }
+  Rmean /= Npts;
+  Dmean /= Npts;
+  Tmean /= Npts;
+
+  float Trange = Tmax - Tmin;
+  // fprintf (stderr, "R,D : %f,%f, T: %f, Trange: %f, Tmin: %f, Tmax: %f\n", Rmean, Dmean, Tmean, Trange, Tmin, Tmax);
+
+  /* project coordinates to a plane centered on the object with units of arcsec */
+  Coords coords;
+  InitCoords (&coords, "DEC--SIN");
+  coords.crval1 = Rmean;
+  coords.crval2 = Dmean;
+  coords.cdelt1 = coords.cdelt2 = 1.0 / 3600.0;
+
+  double *X, *Y, *t, *dX, *dY;
+  ALLOCATE (X, double, N);
+  ALLOCATE (Y, double, N);
+  ALLOCATE (dX, double, N);
+  ALLOCATE (dY, double, N);
+  ALLOCATE (t, double, N);
+
+  int n = 0;
+  for (i = 0; i < N; i++) {
+    if (mask && !mask[i]) continue;
+    RD_to_XY (&X[n], &Y[n], R[i], D[i], &coords);
+    dX[n] = dR[i];
+    dY[n] = dD[i];
+    t[n] = (T[i] - Tmean) / 365.25;
+    n++;
+  }
+
+  PMFit_IRLS fit;
+  if (!FitPMonly_IRLS (&fit, X, dX, Y, dY, t, n, VERBOSE)) {
+    return FALSE;
+  }
+
+  // fprintf (stderr, "Roff, Doff: %f, %f; dRo, dDo: %f, %f\n", fit.Ro, fit.Do, fit.dRo, fit.dDo);
+  
+  XY_to_RD (&Rmean, &Dmean, fit.Ro, fit.Do, &coords);
+  if (VERBOSE) {
+    fprintf (stderr, "Ro, Do: %f, %f +/- %f, %f\n", Rmean, Dmean, fit.dRo, fit.dDo);
+    fprintf (stderr, "uR, uD: %f, %f; duR, duD: %f, %f\n", fit.uR, fit.uD, fit.duR, fit.duD);
+    fprintf (stderr, "chisq: %f Nfit %d\n", fit.chisq, fit.Nfit);
+  }
+
+  set_variable ("RA",   Rmean);
+  set_variable ("DEC",  Dmean);
+  set_variable ("dR",   fit.dRo);
+  set_variable ("dD",   fit.dDo);
+  set_variable ("uR",   fit.uR);
+  set_variable ("uD",   fit.uD);
+  set_variable ("duR",   fit.duR);
+  set_variable ("duD",   fit.duD);
+  set_variable ("plx",  0.0);
+  set_variable ("dplx", 0.0);
+  
+  set_variable ("Tmean",  Tmean);
+  set_variable ("Trange", Trange);
+  set_variable ("Prange", 0.0);
+
+  set_variable ("chisq", fit.chisq);
+  set_variable ("Nfit",  fit.Nfit);
+
+  return (TRUE);
+}
+
+/* do we want an init function which does the alloc and a clear function to free? */
+int FitPMonly_IRLS (PMFit_IRLS *fit, double *X, double *dX, double *Y, double *dY, double *T, int Npts, int VERBOSE) {
+
+  int i;
+
+  static double **A, **B;
+
+
+  double chisq, Xf, Yf;
+
+  double **Cov;
+  double *Beta, *Beta_prev;
+  
+  double sigma_robust, sigma_ols, sigma_final, sigma_hat;
+  double *Wx, *Wy;
+  double *rx, *ry;
+  double *ux, *uy;
+  int dof = 2 * Npts - 4;
+  int p   = 4;
+  int n   = 2 * Npts;
+  double tolerance;
+  int convergence;
+  int iterations;
+  
+  /* do I need to do this as 2 2x2 matrix equations? */
+  if (A == NULL) {
+    ALLOCATE (A, double *, 4);
+    ALLOCATE (B, double *, 4);
+    for (i = 0; i < 4; i++) {
+      ALLOCATE (A[i], double, 4);
+      ALLOCATE (B[i], double, 1);
+      memset (A[i], 0, 4*sizeof(double));
+      memset (B[i], 0, 1*sizeof(double));
+    }
+  }
+
+  // things we need
+  ALLOCATE (Cov, double *, 4);
+  for (i = 0; i < 4; i++) {
+    ALLOCATE ( Cov[i], double, 4);
+  }
+
+  ALLOCATE(Beta, double, 4);
+  ALLOCATE(Beta_prev, double, 4);
+  ALLOCATE(Wx, double, Npts);
+  ALLOCATE(Wy, double, Npts);
+  ALLOCATE(rx,  double, Npts);
+  ALLOCATE(ry,  double, Npts);
+  ALLOCATE(u,  double, Npts);
+  
+  // Convert the measurement errors into initial weights.
+  for (i = 0; i < Npts; i++) {
+    Wx[i] = 1 / dX[i];
+    Wy[i] = 1 / dY[i];
+  }
+  
+  // Solve OLS equation  
+  if (!weighted_LS(T,X,Wx,Y,Wy,Npts,
+		   A,B,VERBOSE)) {
+    // Handle fail case
+    return(FALSE);
+  }
+
+  // Calculate r vector of residuals and least squares sigma
+  sigma_ols = 0.0;
+  for (i = 0; i < Npts; i++) {
+    rx[i] = X[i] - (T[i] * B[0][0] + B[1][0]);
+    ry[i] = Y[i] - (T[i] * B[2][0] + B[3][0]);
+    //    u[i] = r[i] /
+    sigma_ols += SQ(rx[i]) + SQ(ry[i]);
+
+  }
+  sigma_ols = sqrt(sigma_ols / dof);
+
+  // Save OLS covariance;
+  for (i = 0; i < 4; i++) {
+    for (j = 0; j < 4; j++) {
+      Cov[i][j] = A[i][j];
+    }
+  }
+
+  // Save Beta
+  for (i = 0; i < 4; i++) {
+    Beta[i] = B[i][0];
+  }
+
+  // Iterately reweight and solve
+  converged = FALSE;
+  iterations = 0;
+  do {
+    // Save Beta.
+    for (i = 0; i < 4; i ++) {
+      Beta_prev[i] = Beta[i];
+    }
+
+    // Assign W
+    for (i = 0; i < Npts; i++) {
+      Wx[i] = weight_cauchy(rx[i] / dX[i]);
+      Wy[i] = weight_cauchy(ry[i] / dY[i]);
+    }    
+
+    // Solve
+    if (!weighted_LS(T,X,Wx,Y,Wy,Npts,
+		     A,B,VERBOSE)) {
+      // Handle fail case
+      return(FALSE);
+    }
+
+    for (i = 0; i < 4; i++) {
+      Beta[i] = B[i][0];
+    }
+
+    // r
+    sigma_hat = 0.0;
+    for (i = 0; i < Npts; i++) {
+      rx[i] = X[i] - (T[i] * B[0][0] + B[1][0]);
+      ry[i] = Y[i] - (T[i] * B[2][0] + B[3][0]);
+      u[i] = sqrt(SQ(rx[i] / dX[i]) + SQ(ry[i] / dY[i]));
+    }
+    sigma_hat = MAD(u,Npts) / 0.6745;
+    
+    // Check convergence
+    converged = TRUE;
+    tolerance = 1e-4;  // This should probably be tunable.
+    for (i = 0; i < 4; i++) {
+      if (fabs(Beta[i] - Beta_prev[i]) > tolerance * abs(Beta[i])) {
+	converged = FALSE;
+      }
+    }
+
+    iterations++;
+    if (iterations >= 10) {
+      converged = TRUE;
+      // Throw a warning or something here.
+    }
+    
+  } while (!converged);
+
+  double ax, ay;
+  double bx, by;
+  double lambda;
+  double sigma_robust_x, sigma_robust_y;
+  double sigma_final_x,  sigma_final_y;
+  double Sum_Wx, Sum_Wy;
+  
+  ax = 0.0; ay = 0.0;
+  bx = 0.0; by = 0.0;
+  lambda = 0.0;
+  for (i = 0; i < Npts; i++) {
+    Wx[i] = weight_cauchy(rx[i] / dX[i]);
+    Wy[i] = weight_cauchy(ry[i] / dY[i]);
+    
+    ax += dpsi_cauchy(rx[i] / dX[i]);
+    ay += dpsi_cauchy(ry[i] / dY[i]);
+
+    bx += SQ(Wx[i]);
+    by += SQ(Wy[i]);
+
+    Sum_Wx += Wx[i];
+    Sum_Wy += Wy[i];
+  }
+  ax /= 1.0 * Npts;  // mean(psi_dot(r))
+  ay /= 1.0 * Npts; 
+  bx /= 1.0 * (Npts - p); // mean(psi^2(r)) * (N / (N-p))
+  by /= 1.0 * (Npts - p);
+  
+  sigma_robust_x = lamba * sqrt(bx) * sigma_hat * 2.385 / ax;
+  sigma_robust_y = lamba * sqrt(by) * sigma_hat * 2.385 / ay;
+
+  // This is actually sigma^2, as that's the factor in the covariance (dumouchel 4.1)
+  sigma_final_x  = fmax(SQ(sigma_robust_x), (n * SQ(sigma_robust_x) + SQ(p * sigma_OLS)) / (n + SQ(p)));
+  sigma_final_y  = fmax(SQ(sigma_robust_y), (n * SQ(sigma_robust_y) + SQ(p * sigma_OLS)) / (n + SQ(p)));
+
+  for (i = 0; i < 4; i++) {
+    for (j = 0; j < 4; j++) {
+      // This uses the original OLS covariance.
+      if ((i < 2)&&(j < 2)) { // Upper portion
+	Cov[i][j] *= sigma_final_x;
+      }
+      elsif ((i > 1)&&(j > 1)) { // Lower portion
+	Cov[i][j] *= sigma_final_y;
+      }
+      else { // Cross term
+	Cov[i][j] *= sqrt(sigma_final_x * sigma_final_y);
+      }
+    }
+  }
+
+  // Finish.
+  fit[0].Ro = Beta[0];
+  fit[0].uR = Beta[1];
+  fit[0].Do = Beta[2];
+  fit[0].uD = Beta[3];
+  
+  fit[0].dRo = sqrt(Cov[0][0]);
+  fit[0].duR = sqrt(Cov[1][1]);
+  fit[0].dDo = sqrt(Cov[2][2]);
+  fit[0].duD = sqrt(Cov[3][3]);
+
+  // Sort out the final weight threshold.
+
+  // add up the chi square for the fit
+  chisq = 0.0;
+  fit[0].Nfit = 0;
+  for (i = 0; i < Npts; i++) {
+    if ((Wx[i] > 0.1 * Sum_Wx / (1.0 * Npts))||
+	(Wy[i] > 0.1 * Sum_Wy / (1.0 * Npts))) {
+      Xf = fit[0].Ro + fit[0].uR*T[i];
+      Yf = fit[0].Do + fit[0].uD*T[i];
+      chisq += SQ(X[i] - Xf) / SQ(dX[i]);
+      chisq += SQ(Y[i] - Yf) / SQ(dY[i]);
+      fit[0].Nfit += 1;
+    }
+    // if (VERBOSE) fprintf (stderr, "chisq contrib : %f %f : %f %f : %f %f : %f %f : %f\n", Xf, Yf, X[i] - Xf, Y[i] - Yf, dX[i], dY[i], (X[i] - Xf) / dX[i], (Y[i] - Yf) / dY[i], chisq);
+  }
+  //  fit[0].Nfit = Npts;
+
+  // the reduced chisq is divided by (Ndof = 2*Npts - 4)
+  fit[0].chisq = chisq / (2.0*Npts - 4.0);
+  return (TRUE);
+}
+
+
+double weight_cauchy (double x) {
+  double r = x / 2.385;
+  return (1.0 / (1.0 + SQ(r)));
+}
+
+// dpsi = (d/dx) (x * weight(x))
+double dpsi_cauchy (double x) {
+  double r2 = SQ(x / 2.385);
+  return ((1.0 - r2) / (SQ(1 + r2)));
+}
+
+
+// median absolute deviation
+// MAD = median(abs(x - median(x)))
+double MAD(double *in, int N) {
+  double *x;
+  double median = 0.0;
+  int i;
+  
+  ALLOCATE(x,double,N);
+  for (i = 0; i < N; i++) {
+    x[i] = in[i];
+  }
+
+  dsort(x,N);
+
+  if (N % 2) {
+    median = 0.5*(x[(int)(0.5*N)] + x[(int)(0.5*N) - 1]);
+  } else {
+    median = x[(int)(0.5*N)];
+  }
+
+  for (i = 0; i < N; i++ ) {
+    x[i] = fabs(x[i] - median);
+  }
+
+  dsort(x,N);
+
+  if (N % 2) {
+    median = 0.5*(x[(int)(0.5*N)] + x[(int)(0.5*N) - 1]);
+  } else {
+    median = x[(int)(0.5*N)];
+  }
+
+  return(median);
+}
+    
+  
+  
+int weighted_LS (double *T, double *X, double *WX, double *Y, double *WY, int Npts,
+		 double **A, double **B, int VERBOSE) {
+
+  int i,j;
+  double Wx, Wy, Tx, Ty, Tx2, Ty2, Xs, Ys, XT, YT;
+  Wx = Wy = Tx = Ty = Tx2 = Ty2 = Xs = Ys = XT = YT = 0.0;
+  for (i = 0; i < Npts; i++) {
+    Wx += WX[i];
+    Wy += WY[i];
+
+    Tx += T[i]*WX[i];
+    Ty += T[i]*WY[i];
+    
+    Tx2 += SQ(T[i])*WX[i];
+    Ty2 += SQ(T[i])*WY[i];
+    
+    Xs += X[i]*WX[i];
+    Ys += Y[i]*WY[i];
+
+    XT += X[i]*T[i]*WX[i];
+    YT += Y[i]*T[i]*WY[i];
+  }
+
+  // X^T W X
+  A[0][0] = Wx;
+  A[0][1] = Tx;
+
+  A[1][0] = Tx;
+  A[1][1] = Tx2;
+
+  A[2][2] = Wy;
+  A[2][3] = Ty;
+
+  A[3][2] = Ty;
+  A[3][3] = Ty2;
+
+  // X^T W Y
+  B[0][0] = Xs;
+  B[1][0] = XT;
+  B[2][0] = Ys;
+  B[3][0] = YT;
+
+  if (!dgaussjordan ((double **)A, (double **)B, 4, 1)) {
+    if (VERBOSE) fprintf (stderr, "error in fit\n");
+    if (VERBOSE == 2) {
+      int j;
+      for (i = 0; i < 4; i++) {
+	for (j = 0; j < 4; j++) {
+	  fprintf (stderr, "%e ", A[i][j]);
+	}
+	fprintf (stderr, " : %e\n", A[i][0]);
+      }
+    }
+    return FALSE;
+  }
+
+  // A => (X^T W X)^{-1}
+  // B => beta
+  
+  return TRUE;
+}
