agd.Eikonal.HFM_CUDA.SellingAnisotropicDiffusion

  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 numpy as np
  5import cupy as cp
  6import os
  7import numbers
  8
  9from . import cupy_module_helper
 10from .cupy_module_helper import SetModuleConstant
 11from ... import AutomaticDifferentiation as ad
 12
 13def anisotropic_diffusion_scheme(D,dx=1.,ced=None,periodic=False,block_size=256,traits=None):
 14	"""
 15	First stage of solving D_t u = div( D grad u ), with Neumann b.c.
 16	Input : 
 17	- D : positive definite diffusion tensors. Format : (n1,...,nd,d*(d+1)/2)
 18	- ced (optional) : a dictionary with keys
 19	  	- alpha, gamma (real) : weickerts coherence enhancing diffusion parameters
 20	  	- cond_max,cond_amplification_threshold (real, optional) : additional parameters
 21	  	- retD (bool, optional) : wether to return the diffusion tensors 
 22
 23	  	If ced is present, then D is regarded as a structure tensor, and modified 
 24		according to Weickert's coherence enhancing diffusion principle
 25	- dx (optional): grid scale
 26	
 27	Output : 
 28	- dt_max : the largest stable time step for the scheme
 29	- scheme_data : data to be used in anisotropic_diffusion_steps 
 30	- D (optional) : the diffusion tensors, returned if ced['retD']=true 
 31	"""
 32	if D.shape[0]==D.shape[1]==D.ndim-2: # Also allow common format (d,d,n1,...,nd)  
 33		from ... import Metrics
 34		D = np.moveaxis(Metrics.misc.flatten_symmetric_matrix(D),0,-1)
 35	D = cp.ascontiguousarray(D)
 36	shape = D.shape[:-1]
 37	ndim = len(shape)
 38	size = np.prod(shape)
 39	symdim = (ndim*(ndim+1))//2
 40	decompdim = symdim # changes in dimension 4
 41	assert ndim in (1,2,3) and D.shape[-1]==symdim
 42	int_t = np.int32
 43	float_t = np.float32
 44	assert D.dtype==float_t
 45
 46	traits_default = {
 47		'ndim_macro':ndim,
 48		'Int':int_t,
 49		'Scalar':float_t,
 50		'prev_coef':1, # Set prev_coef = 0 for the pure linear operator
 51		'fourth_order_macro':False,
 52	}
 53	if traits is not None: traits_default.update(traits)
 54	traits = traits_default
 55
 56	if isinstance(periodic,numbers.Number): periodic = (periodic,)*ndim
 57	if any(periodic): 
 58		traits['periodic_macro'] = 1
 59		traits['periodic_axes'] = periodic
 60	if ced is not None: 
 61		assert ndim>1
 62		ced = ced.copy()
 63		traits['ced_macro']=1
 64		traits['retD_macro'] = ced.pop('retD',False)
 65	retD = (np.empty_like(D),) if traits.get('retD_macro',False) else tuple()
 66
 67	cuda_path = os.path.join(os.path.dirname(os.path.realpath(__file__)),"cuda")
 68	date_modified = cupy_module_helper.getmtime_max(cuda_path)
 69	source = cupy_module_helper.traits_header(traits,size_of_shape=True)
 70
 71	source += [
 72	'#include "Kernel_SellingAnisotropicDiffusion.h"',
 73	f"// Date cuda code last modified : {date_modified}"]
 74	cuoptions = ("-default-device", f"-I {cuda_path}") 
 75
 76	source = "\n".join(source)
 77	module = cupy_module_helper.GetModule(source,cuoptions)
 78	SetModuleConstant(module,'shape_tot',shape,int_t)
 79	SetModuleConstant(module,'size_tot',size,int_t)
 80	if isinstance(dx,numbers.Number): dx = (dx,)*ndim
 81	SetModuleConstant(module,'dx',dx,float_t)
 82	if ced is not None:
 83		SetModuleConstant(module,"ced_alpha",ced.pop("α"),float_t)
 84		SetModuleConstant(module,"ced_gamma",ced.pop("γ"),float_t)
 85		SetModuleConstant(module,"ced_cond_max",ced.pop("cond_max",100),float_t)
 86		SetModuleConstant(module,"ced_cond_amplification_threshold",
 87			ced.pop("cond_amplification_threshold",2),float_t)
 88		if ced: print(f"Warning : unused ced keys {list(ced.keys())}")
 89
 90	cupy_kernel = module.get_function("anisotropic_diffusion_scheme")
 91	wdiag  = np.full_like(D,0.,shape=shape)
 92	wneigh = np.empty_like(D,shape=shape+(decompdim,));
 93	nneigh = 4 if traits['fourth_order_macro'] else 2
 94	ineigh = np.full_like(D,-1,shape=shape+(decompdim,nneigh),dtype=int_t);
 95
 96	grid_size = int(np.ceil(size/block_size))
 97	cupy_kernel((grid_size,),(block_size,),(D,wdiag,wneigh,ineigh)+retD)
 98
 99	dt_max = 2./np.max(wdiag)
100	return (dt_max,(wdiag,wneigh,ineigh,module))+retD
101
102def anisotropic_diffusion_steps(u,dt,ndt,scheme_data,
103	block_size=1024,overwrite_u=False,out=None):
104	"""
105	Solving D_t u = div( D grad u ), with Neumann b.c.
106	Input : 
107	- u : initial data
108	- dt : time step
109	- ndt : number of time steps
110	- scheme_data : output of anisotropic_diffusion_scheme 
111	"""
112	if not overwrite_u: u=u.copy()
113	u = cp.ascontiguousarray(u)
114	if out is None: out = np.empty_like(u)
115	wdiag,wneigh,ineigh,module = scheme_data
116	float_t = wneigh.dtype 
117	assert u.dtype==float_t
118	assert u.shape==scheme_data[0].shape
119	assert out.shape==u.shape and out.dtype==u.dtype 
120	assert u.flags['C_CONTIGUOUS'] and out.flags['C_CONTIGUOUS']
121	SetModuleConstant(module,'dt',dt,float_t)
122	cupy_kernel = module.get_function("anisotropic_diffusion_step")
123	uold,unew = u,out
124	grid_size = int(np.ceil(u.size/block_size))
125	for i in range(ndt):
126		unew[:]=0.
127		cupy_kernel((grid_size,),(block_size,),(uold,unew,wdiag,wneigh,ineigh))
128		uold,unew = unew,uold
129	if out is not uold: out[:]=uold # unew and uold were just swapped
130	return out