/* 	extrap.java
	
	Wavefield extrapolation program for monochromatic waves.
	Computes a number of x-z timeslices, displayed as a movie.
	Adapted from Jon Claerbout's Imaging the Earth's Interior (1985), p. 105,
	by John N. Louie, 28 November 1998.
*/

/* Include the usual classes */
import java.util.*;
import java.io.*;
import java.lang.Math;

/* Declare a class with a main() method to contain the finite-difference code. */
/* When you change the code in this file, make a copy to a new file with
    a new name like "ex1extrap.java", and be sure to make the same change
    to the class name below - e.g. "class ex1extrap" */
class extrap
	{
	/* The main() method compiles into a file named "extrap.class" that
	    you run in the Java virtual machine as a stand-alone application
	    with a command such as "java extrap"; or however you have
	    renamed it. */
	public static void main(String args[]) /* No runtime arguments needed */
		{
		System.out.println("Starting extrap...");
		/* Declare variables */
		int ix, nx, iz, nz, iw, nw, it, nt;
		float phase, pi2, dx, dz, v;
		float z0, x0, dt, dw, lambda, w, wov, x;
		Complex aa = new Complex();
		Complex a = new Complex();
		Complex b = new Complex();
		Complex c = new Complex();
		Complex cshift = new Complex();
		Complex num = new Complex();
		Complex den = new Complex();
		Complex num2 = new Complex();
		Complex bl = new Complex();
		Complex br = new Complex();
		Complex qm1 = new Complex();
		Complex qnx = new Complex();
		Complex star[][];
		/* We let star[][] be a simple 2-d array of class Complex objects,
		   but it is easier to let p be a class FltVol object from the
		   JRG package, and cd, ce, cf, and q to be class ComplexVec. */
		FltVol p;
		ComplexVec cd, ce, cf, q;
		/* Set dimensions */
		nt = 12;
		nx = 48;
		nz = 96;
		dx = 2;
		dz = 1;
		/* Dimension arrays */
		p = new FltVol(nt, nx, nz);
		cd = new ComplexVec(nx);
		ce = new ComplexVec(nx);
		cf = new ComplexVec(nx);
		q = new ComplexVec(nx);
		star = new Complex[2][3];
		/* Still have to actually create the class Complex
		   objects referenced by the star 2-d array. */
		
		/* 2 frequencies */
		nw = 2;
		/* Initialize constants */
		pi2 = (float)(2.0*Math.PI);
		v = 1;
		lambda = nz*dz/4F;
		dw = v*pi2/lambda;
		dt = pi2/(nt*dw);
		/* Boundary Conditions */
		/* Zero-slope boundary */
		bl.re = 1.0F;
		bl.im = 0.0F;
		br.re = 1.0F;
		br.im = 0.0F;
		
		/* Superimpose nw frequencies */
		for (iw=0; iw<nw; iw++)
			{
			w = (iw+1)*dw;
			/* Frequency / velocity */
			wov = w/v;
			x0 = nx*dx/3F;
			z0 = nz*dz/3F;
			/* Initial conditions for a collapsing spherical wave */
			for (ix=0; ix<nx; ix++)
				{
				x = (ix+1)*dx - x0;
				phase = - wov*(float)(Math.sqrt((double)(z0*z0 + x*x)));
				/* q[ix] = e^(i phase) */
				q.vec[ix].re = (float)(Math.cos((double)(phase)));
				q.vec[ix].im = (float)(Math.sin((double)(phase)));
				}

			/* 15-degree approximation difference star -
			   note we still have to instantiate the class Complex
			   members of the star[][] 2-d array. */
			/* aa = dz/(-4 i wov dx^2) */
			aa.re = 0F;
			aa.im = dz/4F/wov/dx/dx;
			star[0][0] = new Complex(-aa.re, -aa.im);
			star[0][1] = new Complex(-1.0F + 2.0F*aa.re, 2.0F*aa.im);
			star[0][2] = new Complex(-aa.re, -aa.im);
			star[1][0] = new Complex(-aa.re, -aa.im);
			star[1][1] = new Complex(1.0F + 2.0F*aa.re, 2.0F*aa.im);
			star[1][2] = new Complex(-aa.re, -aa.im);
			
			/* Extrapolation in depth */
			for (iz=0; iz<nz; iz++)
				{
				/* Diffraction term */
				for (ix=1; ix<nx-1; ix++)
					{
					/* Apply finite difference star */
					/* cd[ix]=-star[0][2]*q[ix+1]-star[0][1]*q[ix]
					   -star[0][0]*q[ix-1] */
					/* Can only put one complex multiplication
					   method call on each line */
					num.mul(star[0][0],q.vec[ix-1]);
					num2.mul(star[0][1], q.vec[ix]);
					num2.add(num, num2);
					cd.vec[ix].mul(star[0][2],q.vec[ix+1]);
					cd.vec[ix].add(cd.vec[ix],num2);
					cd.vec[ix].scale(-1.0F);
					}
				/* Set side conditions */
				qm1.mul(bl,q.vec[0]);
				qnx.mul(br,q.vec[nx-1]);
				/* cd[0]=-star[0][2]*q[1]-star[0][1]*q[0]
				   -star[0][0]*q[-1] */
				/* Can only put one complex multiplication
				   macro on each line */
				num.mul(star[0][0],qm1);
				num2.mul(star[0][1], q.vec[0]);
				num2.add(num2,num);
				cd.vec[0].mul(star[0][2],q.vec[1]);
				cd.vec[0].add(cd.vec[0],num2);
				cd.vec[0].scale(-1.0F);
				/* cd[nx-1]=-star[0][2]*q[nx]-star[0][1]*q[nx-1]
				   -star[0][0]*q[nx-2] */
				num.mul(star[0][0],q.vec[nx-2]);
				num2.mul(star[0][1], q.vec[nx-1]);
				num2.add(num2,num);
				cd.vec[nx-1].mul(star[0][2],qnx);
				cd.vec[nx-1].add(cd.vec[nx-1],num2);
				cd.vec[nx-1].scale(-1.0F);
				
				/* Set the left and right ends of the diagonal */
				/* Put endl into a complex buffer */
				num.mul(star[1][0],bl);
				num.add(num,star[1][1]);
				/* Put endr into a complex buffer */
				num2.mul(star[1][2],br);
				num2.add(num2,star[1][1]);
				/* Solve the complex tridiagonal system */
				ctris(nx, num, star[1][0], star[1][1], star[1][2],
					num2, cd, q, ce, cf);

				/* cshift = e^(i wov dz) */
				cshift.re = (float)(Math.cos((double)(wov*dz)));
				cshift.im = (float)(Math.sin((double)(wov*dz)));

				/* Shifting term */
				for (ix=0; ix<nx; ix++)
					{
					q.vec[ix].mul(q.vec[ix],cshift);
					}
		
				/* Evolution in time */
				for (it=0; it<nt; it++)
					{
					/* cshift = e^(-i w it dt) */
					cshift.re = (float)(Math.cos((double)(-w*it*dt)));
					cshift.im = (float)(Math.sin((double)(-w*it*dt)));
					/* Pressure wavefield is real part of complex */
					for (ix=0; ix<nx; ix++)
						{
						num.mul(q.vec[ix],cshift);
						/* Compute 1-d address into FltVol vector from
						   3-d indices. */
						p.vec[it*p.nv*p.ne + ix*p.ne + iz] += num.re;
						}
					}
				}
			}

		/* Plot data volume in new frame, after setting label info */
		float rms = p.rms();
		p.setAmpUnits("Pressure");
		p.setUnitsofDim(1, "Distance X");
		p.setDUnitofDim(1, dx);
		p.setUnitsofDim(0, "Depth Z");
		p.setDUnitofDim(0, dz);
		p.setUnitsofDim(2, "= Time t");
		p.setDUnitofDim(2, dt);
		p.viewinFrame("Lab 4: Monochrome Extrapolation", true, 0.5F,
			RgCtab.vagray, false, 3F*rms);
		
		/* End of extrap's main() application method. */
		}

	/*	ctris.java

		Complex tridiagonal equation solver for the system:

		|l c 0 . . . . . 0| | q[0] |   | d[0] |
		|a b c 0 . . . . .| |  .   |   |  .   |
		|0 a b c 0 . . . .| |  .   |   |  .   |
		|. .` ` `  . . . .| |  .   |   |  .   |
		|. . .` ` `  . . .| |  .   | = |  .   |
		|. . . .` ` `  . .| |  .   |   |  .   |
		|. . . . 0 a b c 0| |  .   |   |  .   | 
		|. . . . . 0 a b c| |  .   |   |  .   |
		|0 . . . . . 0 a r| |q[n-1]|   |d[n-1]|

		Adapted from Claerbout (1985) p. 98-99,
		by John N. Louie, 28 November 1998.

		Arguments:	n	dimension of arrays
				endl	corner complex value l (above)
				a,b,c	diagonals of the complex
					coefficient matrix
				endr	corner complex value r (above)
				d	ComplexVec data vector of n complex elements
				q	ComplexVec solution vector of n complex elements
				e,f	ComplexVec temporary iteration vectors of
					n complex elements
										*/
	static void ctris(int n, Complex endl, Complex a, Complex b, Complex c,
		Complex endr, ComplexVec d, ComplexVec q, ComplexVec e, ComplexVec f)
		{
		int i;
		Complex den = new Complex();
		Complex num = new Complex();
		
		/* Initialize unknowns e and f */
		e.vec[0].div(a,endl);
		e.vec[0].scale(-1F);
		f.vec[0].div(d.vec[0],endl);
		
		/* Recursions of equations 19a and 19b */
		for (i=1; i<n-1; i++)
			{
			den.mul(c,e.vec[i-1]);
			den.add(b,den);
			e.vec[i].div(a,den);
			e.vec[i].scale(-1);
			num.mul(c,f.vec[i-1]);
			num.sub(d.vec[i],num);
			f.vec[i].div(num,den);
			}

		/* Right-hand boundary condition, equation 20 */
		num.mul(c,f.vec[n-2]);
		num.sub(d.vec[n-1],num);
		den.mul(c,e.vec[n-2]);
		den.add(endr,den);
		q.vec[n-1].div(num,den);
		
		/* Recursion for equation 21 */
		for (i=n-2; i>=0; i--)
			{
			num.mul(e.vec[i],q.vec[i+1]);
			num.add(num,f.vec[i]);
			q.vec[i].re = num.re;
			q.vec[i].im = num.im;
			}
		/* End of ctris static method of class extrap */
		}

	/* End of methods in class extrap */
	}