agd.Eikonal.HFM_CUDA.ProjectionTTI

  1# Copyright 2021 Jean-Marie Mirebeau, University Paris-Sud, CNRS, University Paris-Saclay
  2# Distributed WITHOUT ANY WARRANTY. Licensed under the Apache License, Version 2.0, see http://www.apache.org/licenses/LICENSE-2.0
  3
  4"""
  5This module attempts to find the parameters of a TTI norm, given a hooke tensor.
  6This boils down to minimizing a polynomial over a sphere. We follow a naive approach, where
  7Newton methods are started from a large number of points in the sphere, and the best result 
  8is selected. In principle however, SDP relaxation techniques are applicable.
  9"""
 10
 11
 12import os
 13import numpy as np
 14from ...Metrics import misc
 15from ... import Sphere
 16from ... import LinearParallel
 17from ...Metrics import Seismic 
 18from ... import AutomaticDifferentiation as ad
 19
 20def rotation_xy(α,β):
 21	"""
 22	A composition of two rotations, along the x and y axes, with the given angles.
 23	"""
 24	cα,sα,cβ,sβ,z = np.cos(α),np.sin(α),np.cos(β),np.sin(β),np.zeros_like(α)
 25	return ad.array([
 26		[cβ,z,sβ],
 27		[sα*sβ,cα,-cβ*sα],
 28		[-cα*sβ,sα,cα*cβ]])
 29
 30def rotation_TTI(x,vdim=None):
 31	"""Construction of the rotation associated to the parameter x returned by ProjectionTTI"""
 32	xdim = 1 if vdim==2 else len(x)
 33	if xdim==1: r = LinearParallel.rotation(x)
 34	elif xdim==2: r = rotation_xy(*x)
 35	elif xdim==3: r = Sphere.rotation3_from_ball3(x)[0]
 36	return np.moveaxis(r,0,1) # Reverse the rotation, to match the notebook convention
 37
 38
 39def ProjectionTTI(hooke,ret_rot=True,n_newton=10,samples=None,quaternion=False):
 40	"""
 41	Attempts to produce a TTI norm matching the given Hooke tensors.
 42	Inputs : 
 43	- n_newton = number of Newton steps in the optimization.
 44	- samples = seed points in unit sphere (or number of them)
 45	
 46	Output : 
 47	Returns the score (squared projection error), an HexagonalMaterial, and 
 48	if retx==False : a rotation matrix
 49	if retx==True  : the parameters of the optimal rotation
 50	"""
 51	import cupy as cp
 52	from . import cupy_module_helper
 53	# Check and prepare the Hooke tensors
 54	hdim = len(hooke)
 55	assert hooke.shape[1]==hdim 
 56	assert hdim in (3,6)
 57	vdim = {3:2,6:3}[hdim]
 58	shape = hooke.shape[2:]
 59	n_hooke = np.prod(shape,dtype=int)
 60
 61	float_t = np.float32
 62	int_t = np.int32
 63#	if retx==False: hooke_in = hooke
 64	hooke = misc.flatten_symmetric_matrix(hooke.reshape((hdim,hdim,-1)))
 65	hooke = cp.asarray(np.moveaxis(hooke,0,-1).astype(float_t))
 66
 67	# Prepare the samples, used to initialize the Newton method, which finds the frame
 68	xdim = 1 if vdim==2 else (3 if quaternion else 2)
 69	if samples is None: samples = {1:40,2:400,3:400}[xdim] # default number of samples
 70	if np.ndim(samples)==0:
 71		if xdim==1: # Sample the rotation angle
 72			nX = samples
 73			samples = np.linspace(0,np.pi/2,nX,endpoint=False)[np.newaxis] 
 74		elif xdim==2: # Sample two rotation angles, along the x and y axes 
 75			nX = int(np.round(np.sqrt(samples))) # recall : vti is z-invariant
 76			aX = np.linspace(-np.pi/2,np.pi/2,nX,endpoint=False)
 77			samples = np.array(np.meshgrid(aX,aX,indexing='ij')).reshape(2,-1)
 78		elif xdim==3: # Sample a point in the three dimensional sphere
 79			nX = int(np.round((2*samples)**(1/3))) # Turned quaternion, then rotation
 80			aX,dx = np.linspace(-1,1,nX,retstep=True,endpoint=False)
 81			samples = np.array(np.meshgrid(*(dx/2+aX,)*xdim,indexing='ij'))
 82			inside = (samples**2).sum(axis=0)<=1
 83			samples = samples[:,inside]
 84	assert len(samples)==xdim
 85	samples = samples.reshape(xdim,-1).astype(float_t).T
 86	n_samples = len(samples)
 87	samples = cp.asarray(samples,dtype=float_t)
 88
 89	# Setup the cuda kernel
 90	traits = {'Scalar':float_t, 'xdim_macro':xdim } 
 91
 92	cuda_path = os.path.join(os.path.dirname(os.path.realpath(__file__)),"cuda")
 93	date_modified = cupy_module_helper.getmtime_max(cuda_path)
 94	source = cupy_module_helper.traits_header(traits)
 95
 96	source += ['#include "Kernel_ProjectionTTI.h"',
 97	f"// Date cuda code last modified : {date_modified}"]
 98	cuoptions = ("-default-device", f"-I {cuda_path}") 
 99	source = "\n".join(source)
100	module = cupy_module_helper.GetModule(source,cuoptions)
101
102	def SetCst(name,var,dtype):cupy_module_helper.SetModuleConstant(module,name,var,dtype)
103	SetCst('n_newton',n_newton,int_t)
104	SetCst('n_hooke',n_hooke,int_t)
105	SetCst('n_samples',n_samples,int_t)
106
107	ProjTTI = module.get_function('ProjectionTTI')
108
109	# Prepare arguments
110	score = cp.full((n_hooke,),np.nan,dtype=float_t)
111	x_out = cp.full((n_hooke,xdim),np.nan,dtype=float_t)
112	hexa = cp.full((n_hooke,5),np.nan,dtype=float_t)
113
114	hooke,score,x_out,samples,hexa = [cp.ascontiguousarray(e) 
115		for e in (hooke,score,x_out,samples,hexa)]
116	size_i = 64
117	size_o = int(np.ceil(n_hooke/size_i))
118	ProjTTI( (size_o,),(size_i,), (hooke,score,x_out,samples,hexa))
119	del hooke
120
121	# Post processing
122	score.reshape(shape)
123	x_out = x_out.T.reshape((xdim,*shape))
124	hexa = hexa.T.reshape((5,*shape))
125	if xdim==1: x_out = x_out[0]; score*=2
126	return score,hexa,(rotation_TTI(x_out,vdim) if ret_rot else x_out)