agd.Eikonal

 1# Copyright 2020 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"""
 5The Eikonal package embeds CPU and GPU numerical solvers of (generalized) eikonal 
 6equations. These are variants of the fast marching and fast sweeping method, based on 
 7suitable discretizations of the PDE, and written and C++.
 8
 9Please see the illustrative notebooks for detailed usage instructions and examples:
10https://github.com/Mirebeau/AdaptiveGridDiscretizations
11
12Main object : 
13- dictIn : a dictionary-like structure, used to gathers the arguments of the 
14eikonal solver, and eventually call it.
15"""
16
17import numpy as np
18import importlib
19import functools
20
21from .LibraryCall import GetBinaryDir
22from .run_detail import Cache
23from .DictIn import dictIn,dictOut,CenteredLinspace
24
25
26def _VoronoiDecomposition_noAD(arr,mode=None,steps='Both',*args,**kwargs):
27	"""
28	Calls the FileVDQ library to decompose the provided quadratic form(s),
29	as based on Voronoi's first reduction of quadratic forms.
30	- mode : 'cpu' or 'gpu' or 'gpu_transfer'. 
31	 Defaults to VoronoiDecomposition.default_mode if specified, or to gpu/cpu adequately.
32	- args,kwargs : passed to gpu decomposition method
33	"""
34	from ..AutomaticDifferentiation.cupy_generic import cupy_set,cupy_get,from_cupy
35	if mode is None: mode = VoronoiDecomposition.default_mode
36	if mode is None: mode = 'gpu' if from_cupy(arr) else 'cpu'
37	if mode in ('gpu','gpu_transfer'):
38		from .HFM_CUDA.VoronoiDecomposition import VoronoiDecomposition as VD
39		if mode=='gpu_transfer': arr = cupy_set(arr)
40		out = VD(arr,*args,steps=steps,**kwargs)
41		if mode=='gpu_transfer': out = cupy_get(out,iterables=(tuple,))
42		return out
43	elif mode=='cpu':
44		from ..Metrics import misc
45		from . import FileIO
46		bin_dir = GetBinaryDir("FileVDQ",None)
47		dim = arr.shape[0]; shape = arr.shape[2:]
48		if dim==1: # Trivial case
49			from .. import Selling
50			return Selling.Decomposition(arr)
51		arr = arr.reshape( (dim,dim,np.prod(shape,dtype=int)) )
52		arr = np.moveaxis(misc.flatten_symmetric_matrix(arr),0,-1)
53		vdqIn ={'tensors':arr,'steps':steps}
54		vdqOut = FileIO.WriteCallRead(vdqIn, "FileVDQ", bin_dir)
55		weights = np.moveaxis(vdqOut['weights'],-1,0)
56		offsets = np.moveaxis(vdqOut['offsets'],(-1,-2),(0,1)).astype(int)
57		weights,offsets = (e.reshape(e.shape[:depth]+shape) 
58			for (e,depth) in zip((weights,offsets),(1,2)))
59		if steps=='Both': return weights,offsets
60		objective = vdqOut['objective'].reshape(shape)
61		vertex = vdqOut['vertex'].reshape(shape).astype(int)
62		chg = np.moveaxis(vdqOut['chg'],(-1,-2),(0,1)) 
63		chg=chg.reshape((dim,dim)+shape)
64		return chg,vertex,objective,weights,offsets
65	else: raise ValueError(f"VoronoiDecomposition unsupported mode {mode}")
66
67def VoronoiDecomposition(arr,*args,**kwargs):
68	"""
69	Voronoi decomposition of arr, an array of dxd symmetric positive definite matrices, 
70	with shape (d,d,n1,...nk), and possibly with AD.
71	args,kwargs : see _VoronoiDecomposition_noAD 
72	"""
73	from .. import AutomaticDifferentiation as ad
74	from ..Metrics.misc import flatten_symmetric_matrix as fltsym
75	from .. import LinearParallel as lp
76	res = _VoronoiDecomposition_noAD(ad.remove_ad(arr),*args,**kwargs)
77	if not ad.is_ad(arr): return res
78	if len(arr)==4: raise ValueError("AD unsupported in dimension 4, sorry")
79	λ,e = res[-2],res[-1]
80	D_flat = arr if kwargs.get('flattened',False) else fltsym(arr)
81	eet_flat = fltsym(lp.outer_self(res[-1]))
82	λ_ad = lp.solve_AV(eet_flat.astype(D_flat.dtype),D_flat) # e is integer valued
83	return *res[:-2],λ_ad,e
84
85VoronoiDecomposition.default_mode = None