fs3d-vec.diderot

#version 1.0
/* fs3d-vec.diderot: function sampler, 3D vector synthetic fields

This is based on fs3d-scl.diderot in this same directory, after much
copy-and-pasting (with fewer comments).  As with that program, the
`-which` option will determine which function is sampled; look for
`("x" == which)` in the code to see the start of the function definitions.

This program's printed output needs to be captured to generate
NRRD header with full orientation info. A vector-valued identity
function is sampled via:

	./fs2d-vec -which ident | unu save -f nrrd -o ident.nrrd
	rm out.nrrd

Note that like in [`fs2d-scl.diderot`](../fs2d), the NRRD header
generated by this program assumes:

1. This program was not compiled with `--double`
2. The program is running on a little-endian machine.
*/

input string which ("which function to sample; one of ident, jac, foo, zrot, ntorus");
input int sz0 ("# samples on fastest axis") = 52;
input int sz1 ("# samples on medium axis") = 51;
input int sz2 ("# samples on slowest axis") = 50;
input real width ("width of edge of cube region sampled") = 4;
input vec3 axis ("axis (non-normalized) of rotation of sampling grid") = [1,1,1];
input real angle ("angle (in degrees) of rotation of sampling grid") = 0;
input vec3 shear ("amount of shear between three axes") = [0,0,0];
input vec3 off ("translation offset, in index space, from origin-centered grid") = [0,0,0];
input vec3 waxis ("axis (non-normalized) of rotation of world space") = [1,-1,1];
input real wangle ("angle (in degrees) of rotation of world space") = 0;
input vec3 vaxis ("axis (non-normalized) of rotation of output vectors") = [-1,1,-1];
input real vangle ("angle (in degrees) of rotation of output vectors") = 0;
input vec3 parm0 ("parameters that functions may use") = [0,0,0];
input vec3 parm1 ("more parameters that functions may use") = [0,0,0];
input vec3 parm2 ("even more parameters that functions may use") = [0,0,0];

function tensor[3,3] q2rot(vec4 qq) = [
   [ qq[0]*qq[0] + qq[1]*qq[1] - qq[2]*qq[2] - qq[3]*qq[3],
     2*(qq[1]*qq[2] - qq[0]*qq[3]),
     2*(qq[1]*qq[3] + qq[0]*qq[2]) ],
   [ 2*(qq[1]*qq[2] + qq[0]*qq[3]),
     qq[0]*qq[0] - qq[1]*qq[1] + qq[2]*qq[2] - qq[3]*qq[3],
     2*(qq[2]*qq[3] - qq[0]*qq[1]) ],
   [ 2*(qq[1]*qq[3] - qq[0]*qq[2]),
     2*(qq[2]*qq[3] + qq[0]*qq[1]),
     qq[0]*qq[0] - qq[1]*qq[1] - qq[2]*qq[2] + qq[3]*qq[3] ]
];

real tht = angle*π/180;
vec3 snax = sin(tht/2)*normalize(axis);
tensor[3,3] rot = q2rot([cos(tht/2), snax[0], snax[1], snax[2]]);
/* copy-and-paste, for world rotation */
real wtht = wangle*π/180;
vec3 wsnax = sin(wtht/2)*normalize(waxis);
tensor[3,3] wrot = q2rot([cos(wtht/2), wsnax[0], wsnax[1], wsnax[2]]);
/* copy-and-paste, for vector rotation */
real vtht = vangle*π/180;
vec3 vsnax = sin(vtht/2)*normalize(vaxis);
tensor[3,3] vrot = q2rot([cos(vtht/2), vsnax[0], vsnax[1], vsnax[2]]);

vec3 spc = width*[1.0/(sz0-1), 1.0/(sz1-1), 1.0/(sz2-1)];
vec3 edge0 = rot•[spc[0], 0, 0];
vec3 edir1 = rot•[0, spc[1], 0];
vec3 edir2 = rot•[0, 0, spc[2]];
vec3 edge1 = edir1 + spc[1]*shear[0]*normalize(edge0);
vec3 edge2 = edir2 + spc[2]*(shear[1]*normalize(edge0) + shear[2]*normalize(edir1));
vec3 offws = [off[0]*spc[0], off[1]*spc[1], off[2]*spc[2]];
vec3 orig = -(edge0*(sz0-1) + edge1*(sz1-1) + edge2*(sz2-1))/2 + offws;

function vec3 func(vec3 pos0) {
   vec3 pos = wrot•pos0;
   real x = pos[0];
   real y = pos[1];
   real z = pos[2];
   vec3 ret = [0,0,0];
   /* parenthesized numbers in comments are the integers
      previously used to identify the case */
   if ("ident" == which) {        // identity (0)
      ret = pos;
   } else if ("jac" == which) { // with Jacobian [parm0, parm1, parm2] (1)
      tensor[3,3] jac = [parm0, parm1, parm2];
      ret = jac•pos;
   } else if ("foo" == which) { // something with predictable derivatives (2)
      // Note: this is used by another example ../tensor2/tensor2.diderot
      ret = [x,2*y,4*z] + x*[0.4,0.2,0.1] + x*x*[0,0.5,0] + y*z*[0,0,1];
   } else if ("zrot" == which) { // rotation around Z, w/ ripple in length (3)
      real ripple = 1 + parm0[0]*cos(8*atan2(y,x));
      ret = [0,0,1]×pos * ripple;
   } else if ("ntorus" == which) { // points away from torus (4)
      real th = atan2(y,x);
      vec3 core = [cos(th),sin(th),0]; // major radius == 1
      vec3 diff = pos - core;
      real rr = |diff|; // minor radius
      real amp = cos(rr*π/2)^2 if rr < 1 else 0;
      ret = amp*diff;
   } else {
      // update "input string which" annotation above as cases are added
      print("Sorry, no function defined for which = ", which, "\n");
   }
   return vrot•ret;
}

strand sample(int idx0, int idx1, int idx2) {
   output vec3 out = [0,0,0];
   update {
      if (0 == idx0 && 0 == idx1 && 0 == idx2) {
         print("NRRD0004\n");
         print("# Complete NRRD file format specification at:\n");
         print("# http://teem.sourceforge.net/nrrd/format.html\n");
         /* NOTE: this assumes we haven't been compiled with --double,
            and there isn't currently a way for the program to learn this
            (which in our experience has not been a problem) */
         print("type: float\n");
         print("dimension: 4\n");
         print("sizes: 3 ", sz0, " ", sz1, " ", sz2, "\n");
         print("kinds: 3-vector space space space\n");
         print("centers: none cell cell cell\n");
         // NOTE: this assumes machine endianness
         print("endian: little\n");
         print("encoding: raw\n");
         print("space dimension: 3\n");
         print("space directions: none (", edge0[0], ",", edge0[1], ",", edge0[2],
                               ") (", edge1[0], ",", edge1[1], ",", edge1[2],
                               ") (", edge2[0], ",", edge2[1], ",", edge2[2], ")\n");
         print("space origin: (", orig[0], ",", orig[1], ",", orig[2], ")\n");
         print("data file: out.nrrd\n");
         print("byte skip: -1\n");
      }
      out = func(orig + idx0*edge0 + idx1*edge1 + idx2*edge2);
      stabilize;
   }
}
initially [ sample(idx0, idx1, idx2)
            | idx2 in 0..(sz2-1),  // slowest axis
              idx1 in 0..(sz1-1),  // medium axis
              idx0 in 0..(sz0-1)]; // fastest axis