agd.Eikonal.HFM_CUDA.Eigh

  1# Copyright 2022 Jean-Marie Mirebeau, ENS Paris-Saclay, 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
  4import os
  5import numpy as np
  6import cupy as cp
  7
  8from . import cupy_module_helper
  9from .cupy_module_helper import SetModuleConstant
 10
 11def eigh(m,quaternion=False,flatsym=None,use_numpy=False):
 12	"""
 13	Computes the eigenvalues of the symmetric matrix m, of dimension d<=3.
 14	The cuda routines are likely more accurate, but they are often slow and memory
 15	intensive for large arrays of small matrices.
 16	Input : 
 17	 - m : array of shape (n1,...,nk,d,d) symmetric w.r.t the last two entries, or of 
 18	  shape (n1,...,nk, d(d+1)/2) if flatsym==True
 19	 - flatsym (bool, default: autodetect): format for the symmetric matrices, see above.
 20	 - quaternion (bool): format specifier for the eigenvectors, see below
 21	 - use_numpy : use the numpy.linalg.eigh routine (calls cuda)
 22	Output :
 23	 - λ : array of shape (n1,...,nk,d). The eigenvalues of m, sorted increasingly.
 24	 - v : the eigenvectors of m, in the following format
 25	   quaternion == false : shape (n1,...,nk,d,d)
 26	   - true : export the eigenvectors compactly using the quaternion format
 27	    (d=3), or omitting the second eigenvector (d=2).
 28	"""
 29	float_t = np.float32
 30	int_t = np.int32
 31	block_size = 1024
 32
 33	if flatsym is None: flatsym = not(m.ndim>=2 and m.shape[-1]==m.shape[-2]<=3)
 34	shape = m.shape[:(-1 if flatsym else -2)]
 35	size = np.prod(shape,dtype=int)
 36	d = int(np.floor(np.sqrt(2*m.shape[-1]))) if flatsym else m.shape[-1]
 37
 38	assert m.shape == ( (*shape,(d*(d+1))//2) if flatsym else (*shape,d,d) )
 39	assert isinstance(quaternion,bool) or quaternion is None
 40	assert 1<=d<=3
 41	vshape = (*shape,{1:1,2:2,3:4}[d]) if quaternion else (*shape,d,d) 
 42
 43	if d==1: # Trivial 1D case
 44		v=np.broadcast_to(cp.ones(1,dtype=float_t),vshape)
 45		λ = m.reshape((*shape,d)).astype(float_t)
 46		return λ,v
 47	if use_numpy:
 48		from ... import Sphere
 49		assert not flatsym
 50		if quaternion is None: return np.linalg.eigvalsh(m)
 51		λ,v = np.linalg.eigh(m)
 52		if not quaternion: return λ,v 
 53		v[np.linalg.det(v)<0,:,-1] *= -1 # Turn v into a direct orthogonal basis
 54		v = np.moveaxis(v,(-2,-1),(0,1))
 55		v = Sphere.sphere1_from_rotation2(v) if d==2 else Sphere.sphere3_from_rotation3(v)
 56		v = np.moveaxis(v,0,-1)
 57		return λ,v
 58	if not flatsym: 
 59		m = cp.stack([m[...,i,j] for i in range(d) for j in range(i+1)], axis=-1)
 60	m = cp.ascontiguousarray(m,dtype=float_t)
 61
 62	traits = {
 63	'Scalar':float_t,
 64	'Int':int_t,
 65	'quaternion_macro':bool(quaternion),
 66	'ndim_macro':d,
 67	}
 68
 69	cuda_path = os.path.join(os.path.dirname(os.path.realpath(__file__)),"cuda")
 70	date_modified = cupy_module_helper.getmtime_max(cuda_path)
 71	source = cupy_module_helper.traits_header(traits)
 72
 73	source += [
 74	'#include "Kernel_Eigh.h"',
 75	f"// Date cuda code last modified : {date_modified}"]
 76	cuoptions = ("-default-device", f"-I {cuda_path}") 
 77
 78	source = "\n".join(source)
 79	module = cupy_module_helper.GetModule(source,cuoptions)
 80	def setcst(*args): SetModuleConstant(module,*args)
 81	setcst('size_tot',size,int_t)
 82
 83	λ = cp.empty((*shape,d),dtype=float_t)
 84	grid_size = int(np.ceil(size/block_size))
 85	if quaternion is None: # Special value to get the eigenvalues alone
 86		f = module.get_function('kernel_eigvalsh')
 87		f((grid_size,),(block_size,),(m,λ))
 88		return λ
 89	
 90	v = cp.empty(vshape,dtype=float_t)
 91	f = module.get_function('kernel_eigh')
 92	f((grid_size,),(block_size,),(m,λ,v))
 93	return λ,v
 94
 95def eigvalsh(m,flatsym=False):
 96	"""
 97	Computes the eigenvalues of the symmetric matrix m, of dimension d<=3.
 98	The cuda routines are likely more accurate, but often slow and memory intensive.
 99	Input : 
100	 - m : array of shape (n1,...,nk,d,d) symmetric w.r.t the last two entries, or of 
101	  shape (n1,...,nk, d(d+1)/2) if flatsym==True
102	 - flatsym (bool): format for the symmetric matrices, see above
103
104	Output :
105	 - λ : array of shape (n1,...,nk,d). The eigenvalues of m, sorted increasingly.
106	"""
107	return eigh(m,flatsym=flatsym,quaternion=None)