/* This file contains the source for my experimental simulator for the
   Schroedinger wave equation of quantum mechanics.  It currently
   simulates the behavior of a single electron confined to a
   1-dimensional, 1-angstrom long space under the influence of a
   time-independent potential energy function which can be changed by
   editing the energy() function and recompiling.  The display shows
   the real and imaginary components of the wave function, as a
   function of the position of the electron, superimposed in two
   different colors, on the top half of the display, and the derived
   probability distribution over the electron's position along the
   bottom of the display superimposed over the potential energy
   function.  A key press shows some stats and a mouse button exits. 

   I first wrote this program in the summer of '95 on Jody's Gateway
   Pentium under Linux.  I've run it on RS6000s at IBM, and on HP
   Snakes and various Suns at the AI lab.  Speed helps a lot!

   -mpf 8/28/96 */

#include <X11/Xlib.h>
#include <X11/Xutil.h>
#include <X11/Xos.h>

#include <stdio.h>
#include <malloc.h>

#include "icon.bitmap"
#define BITMAPDEPTH 1

static Display *display;
static int screen;

/*----------------------------------------------------------------------*/

#include <math.h>

/* Physical constants. */
#define h (6.626e-34)		/* Planck's constant in J-sec */
#define hbar (h/(2*M_PI))	/* in J-sec */
#define me (9.109e-31)		/* Electron rest mass in kg */

/* Parameters of simulation.
 npts=50, dt=5e-23 makes a nice fast display, but a little jaggedy.
 Smoother and slower is npts=200, dt=5e-24. */
static const int npts = 50;	/* Number of points of spatial resolution */
static const double dt = 5e-23;	/* time step, in seconds */
static const double wm = 1e-10;	/* width of sim space in meters = 1 angstrom */
/* Parameters of normal distribution for wavefunction at time t=0. */
static const double mu = .25;	/* How far across the window is the hump. */
static const double sigma = .05; /* How wide is the hump compared to
				    the window. */

/* Some constants derived at init time and cached for efficiency. */
static double dx,		/* Dist. in meters btw. neighboring points. */
  oodx,				/* One over dx. */
  ho2mdt;			/* hbar over (2*mass) times dt */

/* Various arrays for keeping up with the wavefunction state. */
static double *real,*imag;	/* State at current time. */
static double *r2,*i2;		/* State being calc'ed for next time. */
static double *pr,*pi;		/* State at prev. time (for erasing). */
static double *dtvoh;		/* Energy of points times dt over hbar. */

/* For tracking our progress. */
static int cycles = 0;		/* Number of iterations done so far. */
static double t = 0;		/* Total simulated time so far. */

/* Computes a normal (Gaussian) distribution over x, given the
   standard parameters mu (mean) and sigma (sqrt of std. dev.).
   This is used for setting up the electron's initial position
   distribution. */

double normal(x,mu,sigma)
     double x,mu,sigma;
{
  double s2p = sqrt(2*M_PI);
  double d = (x-mu)/sigma;
  return (1/(sigma*s2p))*exp(-0.5*d*d);
}

/* Compute the potential energy, in Joules, as a function of position,
   in meters.  This whole function is just an input parameter to the
   simulation, not part of the simulation itself. */

double energy(double x){
  double dx = x-(wm/2);
  /* Below are some of the different energy functions I have played
     with.  All but one of these should be commented out at any
     given time. */

  /* Negative Gaussian potential well. */
  /*return 5e-15 * (1.0 - 0.25 * normal(x/wm,0.5,0.1));*/

  /* Positive Gaussian potential bump. */
  /*return 5e-15 * 0.25 * normal(x/wm,0.8,0.1);*/

  /* Well where energy drops with inverse distance from center, down
     to a cutoff threshold. */
  /*double adx = (dx<0?-dx:dx);
  double p = 4e-15 - (1e-25 / adx);
  if (p > 0) {
    return p;
  } else {
    return 0;
  }*/

  /* Parabolic well for harmonic oscillator. */
  return dx*dx*1e+6;

  /* Constant, nonzero energy.  Apparent wave rotation rate differs
     from zero-energy case.*/
  /*return -24e-15;*/

  /* Constant, zero energy. */
  /*return 0;*/

  /* A step barrier through which to tunnel. */
  /*if (x < wm/2) {
    return 0;
  } else if (x < wm*.6) {
    return 1.5e-15;
  } else {
    return 0;
  }*/
}

/* Initialize the simulation. */

void init_physics() {
  int i;
  dx = wm/npts;		  /* Calc dist. btw. pts. in meters. */
  oodx = 1/dx;	     /* Cache this calculation to avoid many divides later. */
  ho2mdt = (hbar/(2*me))*dt; /* Similarly. */
  /* Allocate the various arrays. */
  real = (double *)malloc(npts*sizeof(double));
  imag = (double *)malloc(npts*sizeof(double));
  r2 = (double *)malloc(npts*sizeof(double));
  i2 = (double *)malloc(npts*sizeof(double));
  pr = (double *)malloc(npts*sizeof(double));
  pi = (double *)malloc(npts*sizeof(double));
  dtvoh = (double *)malloc(npts*sizeof(double));
  /* Initialize the wavefunction & cache the values of the energy function. */
  for(i=0;i<npts;i++){
    double x = i/((double)npts-1);
    double p = normal(x,mu,sigma); /* Probability of finding the
				      electron here initially. */
    /* Change the coefficient inside the sin/cos to set the
       electron's initial velocity. */
    real[i] = sqrt(p)*cos(0*x*2*M_PI);
    imag[i] = sqrt(p)*sin(0*x*2*M_PI);
    /* Other initial wavefunctions I tried that didn't work so well. */
    /*real[i] = sin(8*x*2*M_PI);
    imag[i] = 0;*/
    /*real[i] = sqrt(p);
      imag[i] = 0;*/
    /* Cache the energy function, times appropriate constants. */
    dtvoh[i] = dt*energy(i*wm/(npts-1))/hbar;
  }
  /* Force endpoints to zero. Seems to yield more stable behavior. */
  real[0] = real[npts-1] = 0;
  imag[0] = imag[npts-1] = 0;
}

/* Below is the core routine.  Given the wavefunction over discrete
   points at time t, compute the new wavefunction for time t + dt,
   where dt is our time increment.  This is done by solving
   Schroedinger's wave equation for dPsi/dt, multiplying by our
   not-actually-infinitesimal dt, and adding the resulting dPsi into
   Psi.  I'm sure this is a pretty naive approximation though. */

void cycle() {
  int i;
  for(i=1;i<npts-1;i++){
    double dr,di;

    /* Estimate the first derivative wrt x on either side of the point.*/
    double dr1 = (real[i]-real[i-1])*oodx;
    double dr2 = (real[i+1]-real[i])*oodx;
    /* Estimate the second derivative wrt x at the point. */
    double ddr = (dr2-dr1)*oodx;

    /* Same thing, but for the imaginary part. */
    double di1 = (imag[i]-imag[i-1])*oodx;
    double di2 = (imag[i+1]-imag[i])*oodx;
    double ddi = (di2-di1)*oodx;

    /* Psi = Psi + dPsi, dPsi found by solving Schroedinger's equation. */
    r2[i] = real[i] - ho2mdt*ddi + dtvoh[i]*imag[i];
    i2[i] = imag[i] + ho2mdt*ddr - dtvoh[i]*real[i];
  }
  /* Copy from temporary i2/r2 arrays back to main array. */
  for(i=1;i<npts-1;i++){
    real[i]=r2[i];imag[i]=i2[i];
  }
  /* Increment time counters. */
  t += dt; cycles++;
}

/*----------------------------------------------------------------------*/
/* Pretty much everything below here is uninteresting graphics and
   user interface BS. */

void main(argc,argv)
     int argc;
     char **argv;
{
  Window win;
  unsigned width, height;	/* window size */
  int x = 0, y = 0;		/* window position */
  unsigned border_width = 4;	/* border four pixels wide */
  unsigned display_width, display_height;
  char *window_name = "Schroedinger Wave Simulator";
  char *icon_name = "schroed";
  Pixmap icon_pixmap;
  XSizeHints size_hints;
  XEvent report;
  GC gc,gce,gcred,gcblue;
  XFontStruct *font_info;
  char *display_name = NULL;
  int i;

  /* connect to X server */
  if ( (display=XOpenDisplay(display_name)) == NULL ) {
    fprintf(stderr,
	    "basicwin: cannot connect to X server %s\n",
	    XDisplayName(display_name));
    exit(-1);
  }

  init_physics();

  /* get screen size from display structure macro */
  screen = DefaultScreen(display);

  display_width = DisplayWidth(display,screen);
  display_height = DisplayHeight(display,screen);

  /* size window with enough room for text */
  width = display_width/3, height = display_height/3;

  /* create opaque window */
  win = XCreateSimpleWindow(display,RootWindow(display,screen),
			    x,y,width,height,border_width,
			    WhitePixel(display,screen),
			    BlackPixel(display,screen));

  /* Create pixmap of depth 1 (bitmap) for icon */
  icon_pixmap = XCreateBitmapFromData(display, win, icon_bitmap_bits,
				      icon_bitmap_width,
				      icon_bitmap_height);

  /* initialize size hint property for window manager */
  size_hints.flags = PPosition | PSize | PMinSize;
  size_hints.x = x;
  size_hints.y = y;
  size_hints.width = width;
  size_hints.height = height;
  size_hints.min_width = 175;
  size_hints.min_height = 125;

  /* set properties for window manager (always before mapping) */
  XSetStandardProperties(display,win,window_name,icon_name,
			 icon_pixmap,argv,argc,&size_hints);

  /* Select event types wanted */
  XSelectInput(display,win, ExposureMask | KeyPressMask |
	       ButtonPressMask | StructureNotifyMask);

  load_font(&font_info);

  /* create GC for text and drawing */
  get_GC(win, &gc, font_info, 1);
  get_GC(win, &gce, font_info, 0);
  get_GC(win, &gcred, font_info, 2);
  get_GC(win, &gcblue, font_info, 3);

  /* Display window */
  XMapWindow(display,win);

  for(i=0;i<npts;i++){
    pr[i]=real[i];pi[i]=imag[i];
  }
  /* get events, use first Expose to display text and graphics; */
  /* ConfigureNotify to indicate a resize; ButtonPress or KeyPress to */
  /* exit */
  while (1) {
    int i;
    XNextEvent(display,&report);
    switch(report.type) {
    case Expose:
      /*printf("Expose!\n");*/
      /* get rid of all other Expose events on the queue */
      while (XCheckTypedEvent(display, Expose, &report));

      draw_graphics(win, gc, width, height, gce, gcred, gcblue);
      for(i=0;i<npts;i++){
	 pr[i]=real[i];pi[i]=imag[i];
      }
      for(i=0;i<128;i++)
	 cycle();
      XClearArea(display,win,0,0,1,1,1);
      break;
    case ConfigureNotify:
      /* window has been resized, change width and height to send to
	 draw_text and draw_graphics in next Expose */
      width = report.xconfigure.width;
      height = report.xconfigure.height;
      XClearArea(display,win,0,0,width,height,1);
      break;
    case KeyPress:
      {
	/* Calculate and print some stats of the wavefunction. */
	double maxr,maxi;
	int amr,ami;
	double maxv = 0; int amv = -1;
	double maxpd = 0; int amp = -1;
	maxr = maxi = maxv = 0;
	for(i=0;i<npts;i++){
	  double pd = real[i]*real[i]+imag[i]*imag[i];
	  if (real[i] > maxr) {maxr=real[i]; amr = i; }
	  if (-real[i] > maxr) {maxr=-real[i]; amr = i; }
	  if (imag[i] > maxi) {maxi=imag[i]; ami = i; }
	  if (-imag[i] > maxi) {maxi=-imag[i]; ami = i; }
	  if (pd > maxpd) {maxpd = pd; amp = -1;}
	}

	printf("Cycle = %d, t = %g.\n",cycles,t);
	printf("     Max R = %g@%d, Max I = %g@%d, Max P = %g\n",
	       maxr,amr,maxi,ami,maxpd);
	/*draw_graphics(win, gc, width, height);
	  for(i=0;i<64;i++) cycle();
	  break;*/
      }
      break;
    case ButtonPress:
      XUnloadFont(display,font_info->fid);
      XFreeGC(display,gc);
      XFreeGC(display,gce);
      XFreeGC(display,gcred);
      XFreeGC(display,gcblue);
      XCloseDisplay(display);
      exit(1);
    default:
      /* all events selected by StructureNotifyMask except */
      /* ConfigureNotify are thrown away here, since nothing is done */
      /* with them */
      break;
    } /* end switch */
  } /* end while */
}

get_GC(Window win, GC *gc, XFontStruct *font_info, int foo) {
  unsigned long valuemask = 0;	/* ignore XGCvalues and use defaults */
  XGCValues values;
  unsigned int line_width = 1;
  int line_style = LineSolid;
  int cap_style = CapButt;
  int join_style = JoinRound;
  int dash_offset = 0;
  static char dash_list[] = {
    12, 24 };
  int list_length = 2;

  /* Create default graphics context */
  *gc = XCreateGC(display,win,valuemask,&values);

  /* specify font */
  XSetFont(display,*gc,font_info->fid);

  {
    XColor sdr,edr;
  if (foo == 1) {
    /* specify black foreground since default may be white on white */
    XSetForeground(display,*gc,WhitePixel(display,screen));
  } else if (foo == 0) {
    XSetForeground(display,*gc,BlackPixel(display,screen));
  } else if (foo == 2) {
    XAllocNamedColor(display,DefaultColormap(display,screen),"cyan",
		&sdr,&edr);
    XSetForeground(display,*gc,edr.pixel);
  } else if (foo == 3) {
    XAllocNamedColor(display,DefaultColormap(display,screen),"yellow",
		&sdr,&edr);
    XSetForeground(display,*gc,edr.pixel);
  }}

  /* set line attributes */
  XSetLineAttributes(display,*gc,line_width,line_style,cap_style,
		     join_style);

  /* set dashes to be line_width in length */
  XSetDashes(display,*gc,dash_offset,dash_list,list_length);
}

load_font(XFontStruct **font_info) {
  char *fontname = "9x15";

  /* Access font */
  if ((*font_info = XLoadQueryFont(display,fontname)) == NULL) {
    fprintf(stderr,"Basic: Cannot open 9x15 font\n");
    exit(-1);
  }
}

draw_graphics(win,gc,window_width,window_height,gce,gcred,gcblue)
     Window win;
     GC gc;
     unsigned window_width, window_height;
     GC gce,gcred,gcblue;
{
  double ght = window_height/2;
  unsigned c1 = ght/2;
  unsigned c2 = window_height;
  int i;
  double maxv = 6.0;
  double mv2 = 15;
  
  for(i=0;i<npts;i++){
    double pd = real[i]*real[i]+imag[i]*imag[i];
    unsigned x = i*window_width/npts;
    int h1 = (int)((real[i]/maxv)*ght*0.5);
    int h2 = (int)((imag[i]/maxv)*ght*0.5);
    int h3 = (int)((pd/(maxv*maxv))*ght);
    /*XDrawLine(display,win,gc, x, c1, x, c1-h1);
    XDrawLine(display,win,gc, x, c2, x, c2-h2);
    XDrawLine(display,win,gc, x, c3, x, c3-h3);*/

    if (i==0) {
      /*Erase old points. */
      XDrawPoint(display,win,gce, x, c1-(int)(ght*0.5*pr[i]/maxv));
      XDrawPoint(display,win,gce, x, c1-(int)(ght*0.5*pi[i]/maxv));
      XDrawPoint(display,win,gce, x,
		 c2-(int)(ght*(pr[i]*pr[i]+pi[i]*pi[i])/(mv2)));

      /* Draw new points. */
      XDrawPoint(display,win,gcred, x, c1-h1);
      XDrawPoint(display,win,gcblue, x, c1-h2);
      XDrawPoint(display,win,gc, x, c2-h3);
    } else {
      XDrawLine(display,win,gce,
		(i-1)*window_width/npts, c1-(int)(ght*0.5*pr[i-1]/maxv),
		x, c1-(int)(ght*0.5*pr[i]/maxv));
      XDrawLine(display,win,gce,
		(i-1)*window_width/npts, c1-(int)(ght*0.5*pi[i-1]/maxv),
		x, c1-(int)(ght*0.5*pi[i]/maxv));
      XDrawLine(display,win,gce,
		(i-1)*window_width/npts,
		c2-(int)(ght*(pr[i-1]*pr[i-1]+pi[i-1]*pi[i-1])/(mv2)),
		x, c2-(int)(ght*(pr[i]*pr[i]+pi[i]*pi[i])/(mv2)));

      XDrawLine(display,win,gcred,
		(i-1)*window_width/npts, c1-(int)(ght*0.5*real[i-1]/maxv),
		x, c1-(int)(ght*0.5*real[i]/maxv));
      XDrawLine(display,win,gcblue,
		(i-1)*window_width/npts, c1-(int)(ght*0.5*imag[i-1]/maxv),
		x, c1-(int)(ght*0.5*imag[i]/maxv));
      XDrawLine(display,win,gc,
		(i-1)*window_width/npts,
		c2-(int)(ght*(real[i-1]*real[i-1]+imag[i-1]*imag[i-1])/(mv2)),
		x, c2-(int)(ght*(real[i]*real[i]+imag[i]*imag[i])/(mv2)));

      

    }

    /*printf("%f\n",dtvoh[i]/dt*1e-18);*/
    XDrawPoint(display,win,gc, x, (int)(c2-ght*dtvoh[i]/dt*1e-20));
  }
  XFlush(display);
}