optiphy/tools/generate_input_photons.py

0001 #!/usr/bin/env python3
0002
0003 from collections import OrderedDict as odict
0004 import argparse, logging, os, json
0005 import numpy as np
0006 from optiphy.ana.sample import sample_trig, sample_normals, sample_reject, sample_linear, sample_linspace, sample_disc
0007 from optiphy.ana.sample import xy_grid_coordinates
0008
0009 log = logging.getLogger(__name__)
0010 np.set_printoptions(linewidth=200, suppress=True, precision=3)
0011
0012 def vnorm(v):
0013     norm = np.sqrt((v*v).sum(axis=1))
0014     norm3 = np.repeat(norm, 3).reshape(-1,3)
0015     v /=  norm3
0016     return v
0017
0018
0019 class InputPhotons:
0020     """
0021     The "WEIGHT" has never been used as such.
0022     The (1,3) sphoton.h slot is used for the iindex integer,
0023     hence set the "WEIGHT" to 0.f in order that the int32
0024     becomes zero.
0025     """
0026
0027     WEIGHT = 0.
0028
0029     DEFAULT_BASE = os.path.expanduser("~/.opticks/InputPhotons")
0030     DTYPE = np.float64
0031
0032     X = np.array( [1., 0., 0.], dtype=DTYPE )
0033     Y = np.array( [0., 1., 0.], dtype=DTYPE )
0034     Z = np.array( [0., 0., 1.], dtype=DTYPE )
0035
0036     POSITION = [0.,0.,0.]
0037     TIME = 0.1
0038     WAVELENGTH  = 440.
0039
0040     @classmethod
0041     def BasePath(cls, name=None):
0042         if name is None:
0043             name = os.environ.get("OPTICKS_INPUT_PHOTON", "RandomSpherical100_f4.npy")
0044
0045         return os.path.join(cls.DEFAULT_BASE, name)
0046
0047     @classmethod
0048     def Path(cls, name, ext=".npy"):
0049         prec = None
0050         if cls.DTYPE == np.float32: prec = "_f4"
0051         if cls.DTYPE == np.float64: prec = "_f8"
0052         return os.path.join(cls.DEFAULT_BASE, "%s%s%s" % (name, prec, ext))
0053
0054
0055     @classmethod
0056     def CubeCorners(cls):
0057         """
0058         :return dir: (8,3) array of normalized direction vectors
0059
0060         000  0   (-1,-1,-1)/sqrt(3)
0061         001  1
0062         010  2
0063         011  3
0064         100  4
0065         101  5
0066         110  6
0067         111  7   (+1,+1,+1)/sqrt(3)
0068         """
0069         v = np.zeros((8, 3), dtype=cls.DTYPE)
0070         for i in range(8): v[i] = list(map(float,[ bool(i & 1), bool(i & 2), bool(i & 4)]))
0071         v = 2.*v - 1.
0072         return vnorm(v)
0073
0074     @classmethod
0075     def GenerateCubeCorners(cls):
0076         direction = cls.CubeCorners()
0077         polarization = vnorm(np.cross(direction, cls.Y))
0078
0079         p = np.zeros( (8, 4, 4), dtype=cls.DTYPE )
0080         n = len(p)
0081         p[:,0,:3] = cls.POSITION + direction  # offset start position by direction vector for easy identification purposes
0082         p[:,0, 3] = cls.TIME*(1. + np.arange(n))
0083         p[:,1,:3] = direction
0084         p[:,1, 3] = cls.WEIGHT
0085         p[:,2,:3] = polarization
0086         p[:,2, 3] = cls.WAVELENGTH
0087         return p
0088
0089     @classmethod
0090     def OutwardsCubeCorners(cls):
0091         direction = cls.CubeCorners()
0092         polarization = vnorm(np.cross(direction, cls.Y))
0093
0094         p = np.zeros( (8, 4, 4), dtype=cls.DTYPE )
0095         n = len(p)
0096         p[:,0,:3] = cls.POSITION + direction  # offset start position by direction vector for easy identification purposes
0097         p[:,0, 3] = cls.TIME*(1. + np.arange(n))
0098         p[:,1,:3] = direction
0099         p[:,1, 3] = cls.WEIGHT
0100         p[:,2,:3] = polarization
0101         p[:,2, 3] = cls.WAVELENGTH
0102         return p
0103
0104
0105     @classmethod
0106     def InwardsCubeCorners(cls, radius):
0107         """
0108         :param radius: of start position
0109         :return p: (8,4,4) array of photons
0110         """
0111         log.info(" radius %s " % radius )
0112         direction = cls.CubeCorners()
0113         polarization = vnorm(np.cross(-direction, cls.Y))
0114
0115         p = np.zeros( (8, 4, 4), dtype=cls.DTYPE )
0116         n = len(p)
0117         p[:,0,:3] = radius*direction
0118         p[:,0, 3] = cls.TIME*(1. + np.arange(n))
0119         p[:,1,:3] = -direction
0120         p[:,1, 3] = cls.WEIGHT
0121         p[:,2,:3] = polarization
0122         p[:,2, 3] = cls.WAVELENGTH
0123         return p
0124
0125     @classmethod
0126     def Axes(cls):
0127         """
0128
0129                Z   -X
0130                |  .
0131                | .
0132                |.
0133        -Y......O------ Y  1
0134               /.
0135              / .
0136             /  .
0137            X   -Z
0138           0
0139         """
0140         v = np.zeros((6, 3), dtype=cls.DTYPE)
0141         v[0] = [1,0,0]
0142         v[1] = [0,1,0]
0143         v[2] = [0,0,1]
0144         v[3] = [-1,0,0]
0145         v[4] = [0,-1,0]
0146         v[5] = [0,0,-1]
0147         return v
0148
0149
0150     @classmethod
0151     def GenerateAxes(cls):
0152         direction = cls.Axes()
0153         polarization = np.zeros((6, 3), dtype=cls.DTYPE)
0154         polarization[:-1] = direction[1:]
0155         polarization[-1] = direction[0]
0156
0157         p = np.zeros( (6, 4, 4), dtype=cls.DTYPE )
0158         n = len(p)
0159         p[:,0,:3] = cls.POSITION
0160         p[:,0, 3] = cls.TIME*(1. + np.arange(n))
0161         p[:,1,:3] = direction
0162         p[:,1, 3] = cls.WEIGHT
0163         p[:,2,:3] = polarization
0164         p[:,2, 3] = cls.WAVELENGTH
0165         return p
0166
0167
0168     @classmethod
0169     def Parallelize1D(cls, p, r, offset=True):
0170         """
0171         :param p: photons array of shape (num, 4, 4)
0172         :param r: repetition number
0173         :return pp:  photons array of shape (r*num, 4, 4)
0174
0175         See parallel_input_photons.py for tests/plotting
0176         """
0177         if r == 0:
0178             return p
0179
0180         o = len(p)          # original number of photons
0181         pp = np.repeat(p, r, axis=0).reshape(-1,r,4,4)  # shape (8,10,4,4)
0182
0183         if offset:
0184             for i in range(o):
0185                 dir = p[i,1,:3]
0186                 pol = p[i,2,:3]   # original polarization, a transverse offset direction vector
0187                 oth = np.cross(pol, dir)
0188                 for j in range(r):
0189                     jj = j - r//2
0190                     pp[i,j,0,:3] = jj*oth
0191
0192         return pp.reshape(-1,4,4)
0193
0194
0195     @classmethod
0196     def Parallelize2D(cls, p, rr, offset=True):
0197         """
0198         :param p: original photons, shaped (o,4,4)
0199         :param [rj,rk]: 2d repeat dimension list
0200         :return pp: shaped (o*rj*rk,4,4)
0201
0202         See parallel_input_photons.py for tests/plotting
0203         """
0204         if len(rr) != 2:
0205             return p
0206
0207         rj, rk = rr[0],rr[1]
0208         o = len(p)          # original number of photons
0209         pp = np.repeat(p, rj*rk, axis=0).reshape(-1,rj,rk,4,4)
0210
0211         if offset:
0212             for i in range(o):
0213                 dir = p[i,1,:3]
0214                 pol = p[i,2,:3]           # original polarization, a transverse offset direction vector
0215                 oth = np.cross(pol, dir)  # other transverse direction perpendicular to pol
0216                 for j in range(rj):
0217                     jj = j - rj//2
0218                     for k in range(rk):
0219                         kk = k - rk//2
0220                         pp[i,j,k,0,:3] = jj*oth + kk*pol
0221
0222         return pp.reshape(-1,4,4)
0223
0224     @classmethod
0225     def GenerateXZ(cls, n, mom, x0lim=[-49.,49.],y0=0.,z0=-99.  ):
0226         """
0227
0228                +-----------------------------------+
0229                |                                   |
0230                |                                   |
0231                |                                   |
0232                |                                   |
0233                |                                   |
0234                |                                   |
0235                |                                   |
0236                |                                   |
0237                |                                   |
0238                |                                   |
0239                |                                   |         Z
0240                |          ^ ^ ^ ^ ^ ^ ^            |         |  Y
0241                |          | | | | | | |            |         | /
0242                |          . . . . . . .            |         |/
0243                +-----------------------------------+         +---> X
0244              -100        -49    0     49          100
0245
0246
0247         """
0248         assert len(x0lim) == 2
0249         if n < 0:
0250             n = -n
0251             xx = sample_linspace(n, x0lim[0], x0lim[1] )
0252         else:
0253             xx = sample_linear(n, x0lim[0], x0lim[1] )
0254
0255         pos = np.zeros((n,3), dtype=cls.DTYPE )
0256         pos[:,0] = xx
0257         pos[:,1] = y0
0258         pos[:,2] = z0
0259
0260         p = np.zeros( (n, 4, 4), dtype=cls.DTYPE )
0261         p[:,0,:3] = pos
0262         p[:,0, 3] = cls.TIME
0263         p[:,1,:3] = mom           # mom : Up:self.Z or Down:-self.Z
0264         p[:,1, 3] = cls.WEIGHT
0265         p[:,2,:3] = cls.Y         # pol
0266         p[:,2, 3] = cls.WAVELENGTH
0267         return p
0268
0269     @classmethod
0270     def GenerateRandomSpherical(cls, n):
0271         """
0272         :param n: number of photons to generate
0273
0274         spherical distribs not carefully checked
0275
0276         The start position is offset by the direction vector for easy identification purposes
0277         so that means the rays will start on a virtual unit sphere and travel radially
0278         outwards from there.
0279
0280         """
0281         spherical = sample_trig(n).T
0282         assert spherical.shape == (n,3)
0283
0284         direction = spherical
0285         polarization = vnorm(np.cross(direction,cls.Y))
0286
0287         p = np.zeros( (n, 4, 4), dtype=cls.DTYPE )
0288         p[:,0,:3] = cls.POSITION + direction
0289         p[:,0, 3] = cls.TIME*(1. + np.arange(n))
0290         p[:,1,:3] = direction
0291         p[:,1, 3] = cls.WEIGHT
0292         p[:,2,:3] = polarization
0293         p[:,2, 3] = cls.WAVELENGTH
0294         return p
0295
0296     @classmethod
0297     def GenerateRandomDisc(cls, n):
0298         spherical = sample_trig(n).T
0299         disc_offset = spherical.copy()
0300         disc_offset[:,0] *= 100.
0301         disc_offset[:,1] *= 100.
0302         disc_offset[:,2] = 0.
0303
0304         p = np.zeros( (n, 4, 4), dtype=cls.DTYPE )
0305         p[:,0,:3] = cls.POSITION + disc_offset
0306         p[:,0, 3] = cls.TIME*(1. + np.arange(n))
0307         p[:,1,:3] = cls.Z
0308         p[:,1, 3] = cls.WEIGHT
0309         p[:,2,:3] = cls.X
0310         p[:,2, 3] = cls.WAVELENGTH
0311         return p
0312
0313     @classmethod
0314     def GenerateUniformDisc(cls, n, radius=100.):
0315         offset = sample_disc(n, dtype=cls.DTYPE)
0316         offset[:,0] *= radius
0317         offset[:,1] *= radius
0318         p = np.zeros( (n, 4, 4), dtype=cls.DTYPE )
0319         p[:,0,:3] = cls.POSITION + offset
0320         p[:,0, 3] = 0.1
0321         p[:,1,:3] = -cls.Z
0322         p[:,1, 3] = cls.WEIGHT
0323         p[:,2,:3] = cls.X
0324         p[:,2, 3] = cls.WAVELENGTH
0325         return p
0326
0327     @classmethod
0328     def GenerateGridXY(cls, n, X=100., Z=1000. ):
0329         sn = int(np.sqrt(n))
0330         offset = xy_grid_coordinates(nx=sn, ny=sn, sx=X, sy=X )
0331         offset[:,2] = Z
0332
0333         p = np.zeros( (n, 4, 4), dtype=cls.DTYPE )
0334         p[:,0,:3] = cls.POSITION + offset
0335         p[:,0, 3] = 0.1
0336         p[:,1,:3] = -cls.Z
0337         p[:,1, 3] = cls.WEIGHT
0338         p[:,2,:3] = cls.X
0339         p[:,2, 3] = cls.WAVELENGTH
0340         return p
0341
0342
0343     @classmethod
0344     def CheckTransverse(cls, direction, polarization, epsilon):
0345         # check elements should all be very close to zero
0346         check1 = np.einsum('ij,ij->i',direction,polarization)
0347         check2 = (direction*polarization).sum(axis=1)
0348         assert np.abs(check1).min() < epsilon
0349         assert np.abs(check2).min() < epsilon
0350
0351     @classmethod
0352     def Check(cls, p):
0353         direction = p[:,1,:3]
0354         polarization = p[:,2,:3]
0355         cls.CheckTransverse( direction, polarization, 1e-6 )
0356
0357     CC = "CubeCorners"
0358     ICC = "InwardsCubeCorners"
0359     RS = "RandomSpherical"
0360     RD = "RandomDisc"
0361     UD = "UniformDisc"
0362     UXZ = "UpXZ"
0363     DXZ = "DownXZ"
0364     RAINXZ = "RainXZ"
0365     GRIDXY = "GridXY"
0366     Z230 = "_Z230"
0367     Z195 = "_Z195"
0368     Z1000 = "_Z1000"
0369     X700 = "_X700"
0370     X1000 = "_X1000"
0371     R500 = "_R500"
0372
0373     NAMES = [CC, CC+"10x10", CC+"100", CC+"100x100", RS+"10", RS+"100", ICC+"17699", ICC+"1", RD+"10", RD+"100", UXZ+"1000", DXZ+"1000" ]
0374     NAMES += [RAINXZ+"100", RAINXZ+"1000", RAINXZ+"100k", RAINXZ+"10k" ]
0375     NAMES += [RAINXZ+Z230+"_100", RAINXZ+Z230+"_1000", RAINXZ+Z230+"_100k", RAINXZ+Z230+"_10k", RAINXZ+Z230+"_1M" ]
0376     NAMES += [RAINXZ+Z195+"_100", RAINXZ+Z195+"_1000", RAINXZ+Z195+"_100k", RAINXZ+Z195+"_10k" ]
0377     NAMES += [RAINXZ+Z230+X700+"_100", RAINXZ+Z230+X700+"_1000", RAINXZ+Z230+X700+"_10k" ]
0378     NAMES += [UD+R500+"_10k"]
0379     NAMES += [GRIDXY+X700+Z230+"_10k", GRIDXY+X1000+Z1000+"_40k"    ]
0380
0381     def generate(self, name, args):
0382         if args.seed > -1:
0383             log.info("seeding with %d " % args.seed)
0384             np.random.seed(args.seed)
0385
0386         meta = dict(seed=args.seed, name=name, creator="input_photons.py")
0387         log.info("generate %s " % name)
0388         if name.startswith(self.RS):
0389             num = int(name[len(self.RS):])
0390             p = self.GenerateRandomSpherical(num)
0391         elif name.startswith(self.UXZ) or name.startswith(self.DXZ):
0392             num = None
0393             mom = None
0394             z0 = None
0395             xl = None
0396             if name.startswith(self.UXZ):
0397                 num = int(name[len(self.UXZ):])  # extract num following prefix
0398                 mom = self.Z
0399                 z0 = -99.
0400                 xl = 49.
0401             elif name.startswith(self.DXZ):
0402                 num = int(name[len(self.DXZ):])
0403                 mom = -self.Z
0404                 z0 = 999.
0405                 xl = 200.
0406             else:
0407                 pass
0408
0409             p = self.GenerateXZ(num, mom, x0lim=[-xl,xl],y0=0,z0=z0 )
0410         elif name.startswith(self.RAINXZ):
0411             d = parsetail(name, prefix=self.RAINXZ)
0412             z0 = 1000. if d['Z'] is None else d['Z']
0413             xl = 250.  if d['X'] is None else d['X']
0414             num = d['N']
0415             mom = -self.Z
0416             p = self.GenerateXZ(-num, mom, x0lim=[-xl,xl],y0=0,z0=z0 )
0417         elif name.startswith(self.UD):
0418             d = parsetail(name, prefix=self.UD)
0419             p = self.GenerateUniformDisc(d["N"], radius=d["R"])
0420         elif name.startswith(self.GRIDXY):
0421             d = parsetail(name, prefix=self.GRIDXY)
0422             p = self.GenerateGridXY(d["N"], X=d["X"], Z=d["Z"])
0423         elif name.startswith(self.RD):
0424             num = int(name[len(self.RD):])
0425             p = self.GenerateRandomDisc(num)
0426         elif name == self.CC:
0427             p = self.GenerateCubeCorners()
0428         elif name.startswith(self.CC):
0429             o = self.OutwardsCubeCorners()
0430             sdim = name[len(self.CC):]
0431             if sdim.find("x") > -1:
0432                 rr = list(map(int, sdim.split("x")))
0433                 p = self.Parallelize2D(o, rr)
0434                 meta["Parallelize2D_rr"] = rr
0435             else:
0436                 r = int(sdim)
0437                 p = self.Parallelize1D(o, r)
0438                 meta["Parallelize1D_r"] = r
0439
0440         elif name.startswith(self.ICC):
0441             sradius = name[len(self.ICC):]
0442             radius = float(sradius)
0443             p = self.InwardsCubeCorners(radius)
0444         else:
0445             log.fatal("no generate method for name %s " %  name)
0446             assert 0
0447
0448         self.Check(p)
0449         meta.update(num=len(p))
0450         return p, meta
0451
0452
0453     def __init__(self, name, args):
0454         InputPhotons.DTYPE = args.dtype
0455         npy_path = self.Path(name, ext=".npy")
0456         json_path = self.Path(name, ext=".json")
0457         generated = False
0458         if os.path.exists(npy_path) and os.path.exists(json_path):
0459             log.info("load %s from %s %s " % (name, npy_path, json_path))
0460             p = np.load(npy_path)
0461             meta = json.load(open(json_path,"r"))
0462         else:
0463             p, meta = self.generate(name, args)
0464             generated = True
0465
0466         self.p = p
0467         self.meta = meta
0468         if generated:
0469             self.save()
0470
0471     name = property(lambda self:self.meta.get("name", "no-name"))
0472
0473     def save(self):
0474         npy_path = self.Path(self.name, ext=".npy")
0475         json_path = self.Path(self.name, ext=".json")
0476         fold = os.path.dirname(npy_path)
0477         if not os.path.isdir(fold):
0478             log.info("creating folder %s " % fold)
0479             os.makedirs(fold)
0480
0481         log.info("save %s to %s and %s " % (self.name, npy_path, json_path))
0482         np.save(npy_path, self.p)
0483         json.dump(self.meta, open(json_path,"w"))
0484
0485     def __repr__(self):
0486         return "\n".join([str(self.meta),".p %s" % self.p.dtype, str(self.p.reshape(-1,16))])
0487
0488
0489 def parsetail(name, prefix=""):
0490     """
0491     """
0492     d = { 'N':None,'X':None, 'Y':None, 'Z':None, 'R':None }
0493     assert name.startswith(prefix)
0494     tail = name[len(prefix):]
0495     elem = np.array(list(filter(None, tail.split("_"))) if tail.find("_") > -1 else [tail])
0496     #print(" name:%s. tail:%s. elem:%s." % (name, tail, str(elem)) )
0497
0498     if len(elem) > 1:
0499         for e in elem[:-1]:
0500            if e[0] in 'XYZR':
0501                d[e[0]] = int(e[1:])
0502
0503     if elem[-1].endswith("k"):
0504         num = int(elem[-1][:-1])*1000
0505     elif elem[-1].endswith("M"):
0506         num = int(elem[-1][:-1])*1000000
0507     else:
0508         num = int(elem[-1])
0509
0510     assert not num is None
0511     d['N'] = num
0512     return d
0513
0514
0515 def test_parsetail():
0516     assert parsetail("RainXZ1k", prefix="RainXZ") == dict(N=1000,X=None,Y=None,Z=None,R=None)
0517     assert parsetail("RainXZ10k", prefix="RainXZ") == dict(N=10000,X=None,Y=None,Z=None,R=None)
0518     assert parsetail("RainXZ1M", prefix="RainXZ") == dict(N=1000000,X=None,Y=None,Z=None,R=None)
0519
0520     #pt = parsetail("RainXZ_Z230_1k", prefix="RainXZ")
0521     #print(pt)
0522
0523     assert parsetail("RainXZ_Z230_1k", prefix="RainXZ") == dict(N=1000,X=None,Y=None,Z=230,R=None)
0524     assert parsetail("RainXZ_Z230_10k", prefix="RainXZ") == dict(N=10000,X=None,Y=None,Z=230,R=None)
0525     assert parsetail("RainXZ_Z230_1M", prefix="RainXZ") == dict(N=1000000,X=None,Y=None,Z=230,R=None)
0526
0527     assert parsetail("RainXZ_Z230_X250_1k", prefix="RainXZ") == dict(N=1000,X=250,Y=None,Z=230,R=None)
0528     assert parsetail("RainXZ_Z230_X500_10k", prefix="RainXZ") == dict(N=10000,X=500,Y=None,Z=230,R=None)
0529     assert parsetail("RainXZ_Z230_X1000_1M", prefix="RainXZ") == dict(N=1000000,X=1000,Y=None,Z=230,R=None)
0530
0531
0532 def test_InwardsCubeCorners17699(ip):
0533     sel = "InwardsCubeCorners17699"
0534     ip0 = ip[sel]
0535     p = ip0.p
0536     m = ip0.meta
0537     r = np.sqrt(np.sum(p[:,0,:3]*p[:,0,:3], axis=1 ))  # radii of start positions
0538
0539
0540 def _dtype_arg(s):
0541     if s == 'float32':
0542         return np.float32
0543     elif s == 'float64':
0544         return np.float64
0545     else:
0546         raise argparse.ArgumentTypeError("dtype must be 'float32' or 'float64'")
0547
0548
0549 def main():
0550     parser = argparse.ArgumentParser()
0551     parser.add_argument( "names", nargs="*", default=InputPhotons.NAMES, help="Name stem of InputPhotons array, default %(default)s" )
0552     parser.add_argument( "--level", default='info', help="logging level, default %(default)s" )
0553     parser.add_argument( "--seed", type=int, default=0, help="seed for np.random.seed() or -1 for non-reproducible generation, default %(default)s" )
0554     parser.add_argument( "--dtype", type=_dtype_arg, default=np.float64, help="type of generated values" )
0555
0556     args = parser.parse_args()
0557
0558     fmt = '[%(asctime)s] p%(process)s {%(pathname)s:%(lineno)d} %(levelname)s - %(message)s'
0559     logging.basicConfig(level=getattr(logging, args.level.upper()), format=fmt)
0560
0561     ip = odict()
0562     for name in args.names:
0563         ip[name] = InputPhotons(name, args)
0564         print(ip[name])
0565
0566
0567 if __name__ == '__main__':
0568     if "PARSETAIL" in os.environ:
0569         test_parsetail()
0570     else:
0571         ip = main()
0572         print("ip odict contains all InputPhotons instances ")
0573         print(ip.keys())