plt.imshow(mag1) #plt.figure() #plt.imshow(rho2) #, origin='lower') meshfile = r"d:\msh3.txt" densfile = r"d:\den3.txt" magfile = r"d:\mag3.txt" graoutfile = r"d:\gra3.dat" magoutfile = r"d:\mag3.dat" graoutfile1 = r"d:\gra3n.dat" magoutfile1 = r"d:\mag3n.dat" area = (-200, 200, -600, 600, 0, 600) #x y z shape = (100, 150, 1) # z y x mesh = PrismMesh(area, shape) mesh.addprop('density', 1000.*rho1.ravel()-2650.0001) mesh.addprop('magnetization', mag1.ravel()) mesh.dump(meshfile, densfile, 'density') #输出网格到磁盘,MeshTools3D可视化 mesh.dump(meshfile, magfile, 'magnetization') #生成核矩阵 kernel=[] narea = (-500, 500,-1000, 1000) #y x nshape = (20, 40) xp, yp, zp = gridder.regular(narea, nshape, z=-1) prisms=[] for p in mesh: prisms.append(p) print('kernel') inc, dec = 30, -4 kernelgz = prism.gz_kernel(xp, yp, zp, prisms) for i, layer in enumerate(mesh.layers()):
plt.figure() plt.imshow(mag1) #plt.figure() #plt.imshow(rho2) #, origin='lower') meshfile = r"d:\msh1.txt" densfile = r"d:\den1.txt" magfile = r"d:\mag1.txt" graoutfile = r"d:\gra1.dat" magoutfile = r"d:\mag1.dat" graoutfile1 = r"d:\gra1n.dat" magoutfile1 = r"d:\mag1n.dat" area = (-100, 100, -750, 750, 0, 700) #x y z shape = (100, 150, 1) # z y x mesh = PrismMesh(area, shape) mesh.addprop('density', 1000. * rho1.ravel() - 2529.99997) mesh.addprop('magnetization', mag1.ravel()) mesh.dump(meshfile, densfile, 'density') mesh.dump(meshfile, magfile, 'magnetization') #输出网格到磁盘,MeshTools3D可视化 # #生成核矩阵 kernel = [] narea = (-500, 500, -1000, 1000) #y x nshape = (20, 40) xp, yp, zp = gridder.regular(narea, nshape, z=-1) prisms = [] for p in mesh: prisms.append(p) print('kernel') inc, dec = 30, -4 kernelgz = prism.gz_kernel(xp, yp, zp, prisms) for i, layer in enumerate(mesh.layers()):
from geoist.inversion.mesh import PrismMesh from geoist.vis import giplt from geoist.inversion.regularization import Smoothness,Damping,TotalVariation from geoist.inversion.pfmodel import SmoothOperator from geoist.inversion.hyper_param import LCurve from geoist.pfm import inv3d meshfile = r"d:\msh.txt" densfile = r"d:\den.txt" #生成场源网格 NS-40km, EW-80km, 500个单元,z方向10层 area = (-20000, 20000, -40000, 40000, 2000, 32000) #NS EW Down shape = (10, 20, 5) #nz ny nx mesh = PrismMesh(area, shape) density=np.zeros(shape) density[3:8,9:12,1:4]=1.0 # z x y mesh.addprop('density', density.ravel()) mesh.dump(meshfile, densfile, 'density') #输出网格到磁盘,MeshTools3D可视化 #生成核矩阵 kernel=[] narea = (-28750, 28750,-48750, 48750) #NS, EW nshape = (30, 40) #NS, EW depthz = [] xp, yp, zp = gridder.regular(narea, nshape, z=-1) for i, layer in enumerate(mesh.layers()): for j, p in enumerate(layer): x1 = mesh.get_layer(i)[j].x1 x2 = mesh.get_layer(i)[j].x2 y1 = mesh.get_layer(i)[j].y1 y2 = mesh.get_layer(i)[j].y2 z1 = mesh.get_layer(i)[j].z1 z2 = mesh.get_layer(i)[j].z2
@author: chens """ # 3rd imports import matplotlib.pyplot as plt import numpy as np #from io import StringIO # local imports from geoist import gridder from geoist.inversion import geometry from geoist.pfm import prism from geoist.inversion.mesh import PrismMesh from geoist.vis import giplt meshfile = r"d:\msh.txt" #StringIO() densfile = r"d:\den.txt" #StringIO() mesh = PrismMesh((0, 10, 0, 20, 0, 5), (5, 2, 2)) mesh.addprop('density', 1000.0 * np.random.rand(20)) mesh.dump(meshfile, densfile, 'density') #print(meshfile.getvalue().strip()) #print(densfile.getvalue().strip()) model = [] for i, layer in enumerate(mesh.layers()): for j, p in enumerate(layer): #print(i,j, p) x1 = mesh.get_layer(i)[j].x1 x2 = mesh.get_layer(i)[j].x2 y1 = mesh.get_layer(i)[j].y1 y2 = mesh.get_layer(i)[j].y2 z1 = mesh.get_layer(i)[j].z1 z2 = mesh.get_layer(i)[j].z2 den = mesh.get_layer(i)[j].props model.append(geometry.Prism(x1, x2, y1, y2, z1, z2, den))
class GravInvAbicModel: def __init__(self, nzyx=[4, 4, 4], smooth_components=None, depth_constraint=None, model_density=None, refer_density=None, weights=None, source_volume=None, smooth_on='m', subtract_mean=True, data_dir='/data/gravity_inversion'): self.gpu_id = 2 self.subtract_mean = subtract_mean self._nz, self._ny, self._nx = nzyx self.smooth_on = smooth_on self.data_dir = data_dir self.gen_model_name() self.nobsx = nzyx[2] self.nobsy = nzyx[1] self.source_volume = source_volume if model_density is None: self._model_density = None else: self._model_density = model_density.ravel() self._smooth_components = smooth_components if smooth_components is None: self._smooth_components = (set(weights.keys()) - set(['depth', 'obs', 'bound', 'refer'])) self.constraints = dict() self.constraints_val = dict() if depth_constraint is None: self.constraints['depth'] = np.ones(np.prod(nzyx)) self.constraints_val['depth'] = None else: self.constraints['depth'] = ( depth_constraint.reshape(-1, 1) * np.ones( (1, self._nx * self._ny))).ravel() self.constraints_val['depth'] = 0 if refer_density is None: self.constraints['refer'] = None self.constraints_val['refer'] = None else: self.constraints['refer'] = np.ones(self._nx * self._ny * self._nz) self.constraints_val['refer'] = refer_density.ravel() self._weights = weights if not 'depth' in self._weights.keys(): self._weights['depth'] = 1.0 self.smop = SmoothOperator() self.kernel_op = None self.abic_val = 0 self.log_total_det_val = 0 self.log_prior_det_val = 0 self.log_obs_det_val = 0 self.min_u_val = 0 self.min_density = -1.0e4 self.max_density = 1.0e4 @property def source_volume(self): return self._source_volume @source_volume.setter def source_volume(self, value): self._source_volume = value self.gen_mesh() def gen_model_name(self): self.model_name = '{}x{}x{}'.format(self._nx, self._ny, self._nz) self.fname = pathlib.Path( self.data_dir) / pathlib.Path(self.model_name + '.h5') @property def weights(self): return self._weights @weights.setter def weights(self, values): self._weights = values if not self.kernel_op is None: self.kernel_op.weights = self._weights @property def smooth_components(self): return self._smooth_components @smooth_components.setter def smooth_components(self, values): self._smooth_components = values @property def refer_density(self): return self.constraints_val['refer'].reshape(self._nz, self._ny, self._nx) @refer_density.setter def refer_density(self, value): self.constraints_val['refer'] = value.ravel() @property def nx(self): return self._nx @nx.setter def nx(self, value): self._nx = value self.nobsx = self._nx self.gen_model_name() if not self.constraints['depth'] is None: self.constraints['depth'] = self.constraints['depth'].reshape( self._nz, -1)[:, 0] * np.ones((1, self._nx * self._ny)) self.constraints['depth'] = self.constraints['depth'].ravel() self.constraints['refer'] = np.ones(self._nx * self._ny * self._nz) @property def ny(self): return self._ny @ny.setter def ny(self, value): self._ny = value self.nobsy = self._ny self.gen_model_name() if not self.constraints['depth'] is None: self.constraints['depth'] = self.constraints['depth'].reshape( self._nz, -1)[:, 0] * np.ones((1, self._nx * self._ny)) self.constraints['depth'] = self.constraints['depth'].ravel() self.constraints['refer'] = np.ones(self._nx * self._ny * self._nz) @property def nz(self): return self._nz @nz.setter def nz(self, value): self._nz = value self.gen_model_name() self.constraints['refer'] = np.ones(self._nx * self._ny * self._nz) print("Warning: nz changed. \nDon't forget setting depth constraints.") @property def model_density(self): return (self._model_density.reshape(self.nz, self.ny, self.nx)) @model_density.setter def model_density(self, value): self._model_density = value.ravel() def gen_mesh(self, height=-1): shape = (self._nz, self._ny, self._nx) self.mesh = PrismMesh(self._source_volume, shape) density = np.ones(shape) * 1.0e3 self.mesh.addprop('density', density.ravel()) # generate obs grid # coordinate: x North-South,y East-West # gridder is in the order: (nx,ny) self.obs_area = (self._source_volume[0] + 0.5 * self.mesh.dims[0], self._source_volume[1] - 0.5 * self.mesh.dims[0], self._source_volume[2] + 0.5 * self.mesh.dims[1], self._source_volume[3] - 0.5 * self.mesh.dims[1]) obs_shape = (self.nobsx, self.nobsy) self.xp, self.yp, self.zp = gridder.regular(self.obs_area, obs_shape, z=height) def _gen_walsh_matrix(self): print('generating walsh_matrix') if os.path.exists(self.fname): with h5py.File(self.fname, mode='r') as f: if not 'walsh_matrix' in f.keys(): have_walsh_matrix = False else: have_walsh_matrix = True else: have_walsh_matrix = False if have_walsh_matrix: return walsh_matrix = walsh.walsh_matrix(self._nx * self._ny * self._nz, normalized=True, ordering='sequence2', nxyz=(self._nx, self._ny, self._nz)) walsh_matrix = walsh_matrix.astype(np.float32) step = self._nx * self._ny * self._nz // 4 components = ['0', '1', '2', '3'] with h5py.File(self.fname, mode='a') as f: fgroup = f.create_group('walsh_matrix') for i in range(4): fgroup.create_dataset(components[i], data=walsh_matrix[i * step:(i + 1) * step, :]) def gen_kernel(self, process=1): def calc_kernel(i): return prism.gz(self.xp[0:1], self.yp[0:1], self.zp[0:1], [self.mesh[i]]) if process > 1: #Winodws MP running has possiblely errors. print('Number of process:', process) with Pool(processes=process) as pool: kernel0 = pool.map(calc_kernel, range(len(self.mesh))) else: kernel0 = [calc_kernel(i) for i in range(len(self.mesh))] self.kernel0 = np.array(kernel0).reshape(self.nz, self.ny, self.nx) self.kernel_op = AbicLSQOperator( self.kernel0, depth_constraint=self.constraints['depth'], smooth_components=self._smooth_components, refer_constraint=self.constraints['refer'], weights=self._weights) def _diagvec(self, vec=None, diag=None): if vec.ndim == 1: return vec * diag else: return vec * diag.reshape(1, -1) @timeit def walsh_transform(self, keys=None): if keys is None: keys = ['kernel'] + list(self.constraints.keys()) + list( self._smooth_components) else: keys = keys if use_gpu > 0: import cupy as cp is_stored = dict() for key in keys: is_stored[key] = False if os.path.exists(self.fname): with h5py.File(self.fname, mode='r') as f: for key in keys: try: if '3' in f[key].keys(): is_stored[key] = True if key == 'depth': res = f['depth'][ 'constraint'][:] - self.constraints['depth'] res = np.linalg.norm(res) / np.linalg.norm( self.constraints['depth']) if res > 1.0e-3: is_stored[key] = False if key == 'kernel': res = f['kernel']['source_volume'][:] - np.array( self.source_volume) res = np.linalg.norm(res) / np.linalg.norm( np.array(self.source_volume)) if res > 1.0e-3: is_stored[key] = False except KeyError: continue self._gen_walsh_matrix() logn = int(np.ceil(np.log2(self._nx * self._ny * self._nz))) norm_walsh = 1. / (np.sqrt(2)**logn) blocks = ['0', '1', '2', '3'] matvec_op = { 'kernel': self.kernel_op.gtoep.matvec, 'depth': lambda x: self._diagvec(x, diag=np.sqrt(self.constraints['depth'])) } for key in self._smooth_components: matvec_op[key] = lambda x: self.smop.derivation( x.reshape(-1, self.nz, self.ny, self.nx ), component=key).reshape(x.shape[0], -1) is_stored['refer'] = True for key in keys: if is_stored[key]: print('walsh transformation of {} already exists.'.format(key)) continue print('performing walsh transformation on {}.'.format(key)) step = self.nx * self.ny * self.nz // 4 if key == 'depth': step = self._nz with h5py.File(self.fname, mode='a') as f: try: del f[key] except KeyError: pass dxyz_group = f.create_group(key) walsh_group = f['walsh_matrix'] for i in range(4): print("\t progress {}/4".format(i)) part_walsh = walsh_group[blocks[i]][:] if key == 'depth': part_walsh = walsh_group[blocks[i]][:self._nz] part_walsh = matvec_op[key](part_walsh) if use_gpu > 0: with cp.cuda.Device(self.gpu_id): res = cp.zeros((step, step)) j = 0 while j * step < part_walsh.shape[1]: tmp_block_gpu = cp.asarray( part_walsh[:, j * step:(j + 1) * step]) res += tmp_block_gpu @ tmp_block_gpu.T j += 1 res = cp.asnumpy(res) if key in self._smooth_components: res[np.abs(res) < 1.0e-1 * norm_walsh] = 0. tmp_block_gpu = None mempool = cp.get_default_memory_pool() pinned_mempool = cp.get_default_pinned_memory_pool( ) mempool.free_all_blocks() pinned_mempool.free_all_blocks() else: res = np.zeros((step, step)) j = 0 while j * step < part_walsh.shape[1]: tmp_block_gpu = np.asarray( part_walsh[:, j * step:(j + 1) * step]) res += tmp_block_gpu @ tmp_block_gpu.T j += 1 if key in self._smooth_components: res[np.abs(res) < 1.0e-1 * norm_walsh] = 0. dxyz_group.create_dataset(blocks[i], data=res) if ('depth' in keys) and (not is_stored['depth']): with h5py.File(self.fname, mode='a') as f: try: del f['depth']['constraint'] except KeyError: pass dxyz_group = f['depth'] dxyz_group.create_dataset('constraint', data=self.constraints['depth']) if ('kernel' in keys) and (not is_stored['kernel']): with h5py.File(self.fname, mode='a') as f: try: del f['kernel']['source_volume'] except KeyError: pass dxyz_group = f['kernel'] dxyz_group.create_dataset('source_volume', data=np.array(self._source_volume)) @property def depth_constraint(self): return (self.constraints['depth'].reshape(self._nz, -1)[:, 0]) @depth_constraint.setter def depth_constraint(self, value): self.constraints['depth'] = (value.reshape(-1, 1) * np.ones( (1, self._nx * self._ny))).ravel() @timeit def forward(self, model_density=None): if model_density is None: model_density = self._model_density else: model_density = model_density.ravel() self.obs_data = self.kernel_op.gtoep.matvec(model_density) def _gen_rhs(self): self.rhs = self._weights['obs'] * self.kernel_op.gtoep.rmatvec( self.obs_data) if 'depth' in self._weights.keys(): v = self.constraints['depth'] * self.constraints_val['refer'] if 'refer' in self._weights.keys(): self.rhs += (self._weights['refer'] * self._weights['depth'] * self.constraints['depth'] * v) if self.smooth_on == 'm-m0': for key in self._smooth_components: tmp2 = v.reshape(-1, self._nz, self._ny, self._nx) tmp2 = self.smop.derivation(tmp2, component=key) tmp2 = self.smop.rderivation(tmp2, component=key) if v.ndim == 1: self.rhs += self._weights[key] * self._weights[ 'depth'] * self.constraints['depth'] * tmp2.ravel() else: self.rhs += self._weights[key] * self._weights[ 'depth'] * self.constraints['depth'] * tmp2.reshape( v.shape[0], -1) @timeit def do_linear_solve(self): self.do_linear_solve_quiet() def do_linear_solve_quiet(self): self._gen_rhs() if self.subtract_mean: sum_obs = np.sum(self.obs_data) tmp_b = np.zeros(len(self.rhs) + 1) #np.zeros(len(self.rhs+1)) #chenshi tmp_b[:-1] = self.rhs tmp_b[-1] = sum_obs tmp_op = AbicLSQOperator2(self.kernel_op) self.solution = spsparse.linalg.cg(tmp_op, tmp_b, tol=1.0e-5)[0] else: self.solution = spsparse.linalg.cg(self.kernel_op, self.rhs, tol=1.0e-5)[0] @timeit def calc_u(self, solved=False, x=None): return self.calc_u_quiet(solved, x) @timeit def calc_min_u(self, solved=False, x=None): return self.calc_u_quiet(solved, x) def calc_u_quiet(self, solved=False, x=None): if x is None: if not solved: self.do_linear_solve_quiet() x = self.solution self.min_u_val = self._weights['obs'] * np.linalg.norm( self.kernel_op.gtoep.matvec(x) - self.obs_data)**2 if ('refer' in self._weights.keys()) and (self.smooth_on == 'm-m0'): v = x - self.constraints_val['refer'] else: v = x if 'depth' in self._weights.keys(): v = np.sqrt(self._weights['depth']) * self.constraints['depth'] * v for key in self._smooth_components: tmp2 = self.smop.derivation(v.reshape(self._nz, self._ny, self._nx), component=key) self.min_u_val += self._weights[key] * np.linalg.norm( tmp2.ravel())**2 if 'refer' in self._weights.keys(): v = x - self.constraints_val['refer'] if 'depth' in self._weights.keys(): v = np.sqrt( self._weights['depth']) * self.constraints['depth'] * v self.min_u_val += self._weights['refer'] * np.linalg.norm(v)**2 return self.min_u_val def jac_u(self, x=None): res = self.kernel_op.matvec(x) - self.rhs return 2. * res def hessp_u(self, x, v): res = self.kernel_op.matvec(v) return 2. * res @timeit def bound_optimize(self, x0=None): density_bounds = Bounds(self.min_density, self.max_density) if x0 is None: x0 = np.zeros( self._nx * self._ny * self._nz) + (self.max_density - self.min_density) / 2. self.bound_solution = minimize( lambda x: self.calc_u_quiet(solved=True, x=x), x0, method='trust-constr', jac=self.jac_u, hessp=self.hessp_u, bounds=density_bounds, ) def lasso_target(self, x): # self.min_u_val = self._weights['obs']*np.linalg.norm(self.kernel_op.gtoep.matvec(x) - self.obs_data)**2 self.min_u_val = self._weights['obs'] * np.linalg.norm( self.kernel_op.gtoep.matvec(x) - self.obs_data) if ('refer' in self._weights.keys()) and (self.smooth_on == 'm-m0'): v = x - self.constraints_val['refer'] else: v = x if 'depth' in self._weights.keys(): v = np.sqrt(self._weights['depth']) * self.constraints['depth'] * v for key in self._smooth_components: tmp2 = self.smop.derivation(v.reshape(self._nz, self._ny, self._nx), component=key) self.min_u_val += self._weights[key] * np.linalg.norm(tmp2.ravel()) if 'refer' in self._weights.keys(): v = x - self.constraints_val['refer'] if 'depth' in self._weights.keys(): v = np.sqrt( self._weights['depth']) * self.constraints['depth'] * v self.min_u_val += self._weights['refer'] * np.linalg.norm(v) return self.min_u_val def lasso_jac(self, x): # jac = self.kernel_op.gtoep.rmatvec(self.kernel_op.gtoep.matvec(x)) - self.kernel_op.gtoep.rmatvec(self.obs_data) # jac = 2.0*self._weights['obs']*jac jac = self.kernel_op.gtoep.rmatvec( self.kernel_op.gtoep.matvec(x)) - self.kernel_op.gtoep.rmatvec( self.obs_data) jac = self._weights['obs'] * jac norm_res = np.linalg.norm(self.obs_data - self.kernel_op.gtoep.matvec(x)) jac = jac / norm_res if 'refer' in self.weights.keys(): v = x - self.constraints_val['refer'] if 'depth' in self._weights.keys(): v = self.constraints['depth'] * v norm_refer = np.linalg.norm(v) jac += self.weights['refer'] * self.weights[ 'depth'] * self.constraints['depth'] * v / norm_refer for key in self.smooth_components: v = x if 'depth' in self._weights.keys(): v = self.constraints['depth'] * v tmp2 = self.smop.derivation(v.reshape(self._nz, self._ny, self._nx), component=key) smooth_norm = np.linalg.norm(tmp2.ravel()) tmp2 = self.smop.rderivation(tmp2, component=key) jac += self.weights[key] * self.weights[ 'depth'] * self.constraints['depth'] * tmp2.ravel( ) / smooth_norm return jac def lasso_hessp(self, x, v): # res = self.kernel_op.gtoep.rmatvec(self.kernel_op.gtoep.matvec(v)) # res = 2.0*self._weights['obs']*res norm_res = np.linalg.norm(self.obs_data - self.kernel_op.gtoep.matvec(x)) res = self.kernel_op.gtoep.rmatvec( self.kernel_op.gtoep.matvec(v)) / norm_res gradient_res = ( self.kernel_op.gtoep.rmatvec(self.kernel_op.gtoep.matvec(x)) - self.kernel_op.gtoep.rmatvec(self.obs_data)) res -= np.dot(gradient_res, v) / norm_res**3 * gradient_res res *= self._weights['obs'] if 'refer' in self.weights.keys(): v2 = x - self.constraints_val['refer'] if 'depth' in self._weights.keys(): v2 = self.constraints['depth'] * v2 norm_refer = np.linalg.norm(v2) res += self.weights['refer'] * self.weights[ 'depth'] / norm_refer * self.constraints['depth'] * v grad_ref = self.constraints['depth'] * v2 res -= (self.weights['refer'] * self.weights['depth'] * (np.dot(v, grad_ref) / norm_refer**3 * grad_ref)) for key in self.smooth_components: v2 = x if 'depth' in self._weights.keys(): v2 = self.constraints['depth'] * v2 tmp2 = self.smop.derivation(v2.reshape(self._nz, self._ny, self._nx), component=key) smooth_norm = np.linalg.norm(tmp2.ravel()) tmp2 = self.smop.rderivation(tmp2, component=key) grad_sm = self.constraints['depth'] * tmp2.ravel() tmp2 = v.reshape(self.nz, self.ny, self.nx) tmp2 = self.smop.derivation(tmp2, component=key) tmp2 = self.smop.rderivation(tmp2, component=key) tmp2 = self.constraints['depth'] * tmp2.ravel() res += (self._weights['depth'] * self._weights[key] / smooth_norm * (tmp2 - np.dot(v, grad_sm) * grad_sm / smooth_norm**2)) return res @timeit def lasso_optimize(self, x0=None): density_bounds = Bounds(self.min_density, self.max_density) if x0 is None: x0 = (np.random.rand(self._nx * self._ny * self._nz) * (self.max_density - self.min_density) / 2. + (self.max_density + self.min_density) / 2.) self.bound_solution = minimize(self.lasso_target, x0, method='trust-constr', jac=self.lasso_jac, hessp=self.lasso_hessp, bounds=density_bounds) def calc_res(self): self.residuals = dict() self.stds = dict() self.residuals['obs'] = np.linalg.norm( self.kernel_op.gtoep.matvec(self.solution) - self.obs_data)**2 self.stds['obs'] = np.std( self.kernel_op.gtoep.matvec(self.solution) - self.obs_data) for key in self._smooth_components: tmp2 = self.solution.reshape(self._nz, self._ny, self._nx) if ('refer' in self.constraints_val.keys()) and (self.smooth_on == 'm-m0'): tmp2 -= self.constraints_val['refer'].reshape( self._nz, self._ny, self._nx) tmp2 = self.smop.derivation(tmp2, component=key) self.residuals[key] = np.linalg.norm(tmp2.ravel())**2 self.stds[key] = np.std(tmp2.ravel()) if 'refer' in self.constraints_val.keys(): self.residuals['refer'] = np.linalg.norm( self.solution.ravel() - self.constraints_val['refer'].ravel())**2 self.stds['refer'] = np.std(self.solution.ravel() - self.constraints_val['refer'].ravel()) def calc_log_prior_total_det_quiet(self): self.log_prior_det_val = 0 self.log_total_det_val = 0 blocks = ['0', '1', '2', '3'] prior_eigs = np.zeros(self._nx * self._ny * self._nz) total_eigs = np.zeros(self._nx * self._ny * self._nz) step = self._nx * self._ny * self._nz // 4 try: depth_weight = self._weights['depth'] except KeyError: depth_weight = 1. with h5py.File(self.fname, mode='r') as f: if 'depth' in self._weights.keys(): depth_walsh = f['depth']['0'][:] for i_b, block in enumerate(blocks): tmp_block = np.zeros((step, step)) for dxyz_name in self._smooth_components: try: dxyz_walsh = f[dxyz_name][block][:].reshape( step // self._nz, self._nz, step // self._nz, self._nz) ein_path = np.einsum_path('mi,xiyj,jn->xmyn', depth_walsh.T, dxyz_walsh, depth_walsh, optimize='optimal')[0] tmp_multi = np.einsum('mi,xiyj,jn->xmyn', depth_walsh.T, dxyz_walsh, depth_walsh, optimize=ein_path) tmp_block += depth_weight * self._weights[ dxyz_name] * tmp_multi.reshape(step, step) except KeyError: pass if 'refer' in self._weights.keys(): tmp_multi_small = depth_walsh.T @ depth_walsh for i in range(step // self._nz): tmp_block[i * self._nz:(i + 1) * self._nz, i * self._nz:(i + 1) * self._nz] += depth_weight * self._weights[ 'refer'] * tmp_multi_small if use_gpu > 0: import cupy as cp with cp.cuda.Device(self.gpu_id): tmp_block_gpu = cp.asarray(tmp_block, dtype=np.float32) eigs = cp.linalg.eigvalsh(tmp_block_gpu) prior_eigs[i_b * step:(i_b + 1) * step] = cp.asnumpy(eigs) self.log_prior_det_val += cp.asnumpy( cp.sum(cp.log(eigs))) tmp_block_gpu = None eigs = None free_gpu() tmp_block += self._weights['obs'] * f['kernel'][block][:] with cp.cuda.Device(self.gpu_id): tmp_block_gpu = cp.asarray(tmp_block, dtype=np.float32) eigs = cp.linalg.eigvalsh(tmp_block_gpu) total_eigs[i_b * step:(i_b + 1) * step] = cp.asnumpy(eigs) self.log_total_det_val += cp.asnumpy( cp.sum(cp.log(eigs))) tmp_block_gpu = None eigs = None free_gpu() else: tmp_block_gpu = np.asarray(tmp_block, dtype=np.float32) eigs = np.linalg.eigvalsh(tmp_block_gpu) prior_eigs[i_b * step:(i_b + 1) * step] = eigs self.log_prior_det_val += np.sum(np.log(eigs)) tmp_block_gpu = None eigs = None tmp_block += self._weights['obs'] * f['kernel'][block][:] tmp_block_gpu = np.asarray(tmp_block, dtype=np.float32) eigs = np.linalg.eigvalsh(tmp_block_gpu) total_eigs[i_b * step:(i_b + 1) * step] = eigs self.log_total_det_val += np.sum(np.log(eigs)) tmp_block_gpu = None eigs = None if use_gpu > 0: self.log_prior_det_val = cp.asnumpy(self.log_prior_det_val) self.log_total_det_val = cp.asnumpy(self.log_total_det_val) else: self.log_prior_det_val = self.log_prior_det_val self.log_total_det_val = self.log_total_det_val self.eigs = {'prior': prior_eigs, 'total': total_eigs} return self.log_prior_det_val, self.log_total_det_val @timeit def calc_log_prior_total_det(self): return self.calc_log_prior_total_det_quiet() def calc_log_obs_det_quiet(self): self.log_obs_det_val = np.log(self._weights['obs']) * len( self.obs_data) return self.log_obs_det_val @timeit def calc_log_obs_det(self): return self.calc_log_obs_det_quiet() @timeit def calc_abic(self): '''-log_prior_det_value+log_total_det-log_obs_det+min_u''' self.calc_log_prior_total_det() self.calc_u() self.calc_log_obs_det() self.abic_val = (self.log_total_det_val + self.min_u_val - self.log_prior_det_val - self.log_obs_det_val) return self.abic_val def calc_abic_quiet(self): '''-log_prior_det_value+log_total_det-log_obs_det+min_u''' self.calc_log_prior_total_det_quiet() self.calc_u_quiet() self.calc_log_obs_det_quiet() self.abic_val = (self.log_total_det_val + self.min_u_val - self.log_prior_det_val - self.log_obs_det_val) return self.abic_val def _abic_optimize_exp(self): #optimize_keys = list(set(self._weights.keys())-set(['depth'])) optimize_keys = list(self._weights.keys()) def abic_target(x): for i, key in enumerate(optimize_keys): self._weights[key] = np.exp(x[i]) return self.calc_abic_quiet() x0 = np.zeros(len(self._weights)) for i, key in enumerate(optimize_keys): x0[i] = np.log(self._weights[key]) self.abic_optimize_summary = minimize(abic_target, x0, method='Nelder-Mead') def _abic_optimize_bound(self): optimize_keys = list(self._weights.keys()) def abic_target(x): for i, key in enumerate(optimize_keys): self._weights[key] = x[i] return self.calc_abic_quiet() x0 = np.zeros(len(self._weights)) for i, key in enumerate(optimize_keys): x0[i] = self._weights[key] weight_constraint = LinearConstraint(np.eye(len(optimize_keys)), 0., np.inf) self.abic_optimize_summary = minimize(abic_target, x0, method='COBYLA', constraints=weight_constraint) @timeit def abic_optimize(self): self._abic_optimize_bound() @timeit def para_grad(self, x): pass def u_bound(self): pass def print_summary(self): print('abic values:{}'.format(self.abic_val)) print('log total det:{}'.format(self.log_total_det_val)) print('log prior det:{}'.format(self.log_prior_det_val)) print('log obs det:{}'.format(self.log_obs_det_val)) print('min u:{}'.format(self.min_u_val)) print('std:', end=' ') print(self.stds) print('1/var:', end=' ') print({k: 1. / v**2 for k, v in self.stds.items()}) print('norms:', end=' ') print(self.residuals)
class GravInvAbicModel: def __init__(self,conf_file=None,**kwargs): self.confs = {'nzyx':[4,4,4], 'gpu_id':2, 'smooth_components':None, 'depth_constraint':None, 'model_density':None, 'refer_densities':None, 'weights':None, 'source_volume':None, 'smooth_on':'m', 'subtract_mean':False, 'optimize_keys':None, 'mode':'walsh', 'data_dir':'/data/gravity_inversion'} confs = dict() if not conf_file is None: with open(conf_file) as f: confs = json.load(f) self.confs = {**self.confs,**confs,**kwargs} self._nz,self._ny,self._nx = self.confs['nzyx'] self.nobsx = self._nx self.nobsy = self._ny self.gen_model_name() self.source_volume = self.confs['source_volume'] if self.confs['model_density'] is None: self._model_density = None else: self._model_density = self.confs['model_density'].ravel() self._smooth_components = self.confs['smooth_components'] if self.confs['smooth_components'] is None: self._smooth_components = (set(self.confs['weights'].keys()) - set(['depth', 'obs', 'bound', 'refers'])) self.constraints = dict() self.constraints_val = dict() if self.confs['depth_constraint'] is None: self.constraints['depth'] = np.ones(np.prod(self.nzyx)) self.constraints_val['depth'] = None else: self.constraints['depth'] = (self.confs['depth_constraint'].reshape(-1,1)*np.ones((1,self._nx*self._ny))).ravel() self.constraints_val['depth'] = 0 if self.confs['refer_densities'] is None: self.constraints['refers'] = None self.constraints_val['refers'] = None else: self.constraints['refers'] = np.ones(self._nx*self._ny*self._nz) self.refer_densities = self.confs['refer_densities'] self.n_refer = len(self.confs['refer_densities']) self._weights = self.confs['weights'] if not 'depth' in self._weights.keys(): self._weights['depth'] = 1.0 self.smop = SmoothOperator() self.kernel_op = None self.abic_val = 0 self.log_total_det_val = 0 self.log_prior_det_val = 0 self.log_obs_det_val = 0 self.min_u_val = 0 self.min_density = -1.0e4 self.max_density = 1.0e4 self.optimize_log = {'parameters':[],'abic_vals':[]} def gen_model_name(self): '''generate a file name to save data of current model. The model will be saved in ``self.data_dir`` directory. ''' self.model_name = '{}x{}x{}'.format(self._nx,self._ny,self._nz) self.fname = pathlib.Path(self.data_dir)/pathlib.Path(self.model_name+'.h5') @property def gpuid(self): ''' Which gpu card will be used. Ignored if set ``use_gpu=0``. ''' return self.confs['gpuid'] @gpuid.setter def gpuid(self,values): self.confs['gpuid'] = values @property def smooth_on(self): '''Which variable should be smoothed. Only ``'m'`` is supported right now, which means smooth on density. ''' return self.confs['smooth_on'] @smooth_on.setter def smooth_on(self,values): self.confs['smooth_on'] = values @property def mode(self): '''How to calculate determinants. Could be 'naive' or 'walsh'. ''' return self.confs['mode'] @mode.setter def mode(self,values): self.confs['mode'] = values @property def data_dir(self): '''Specify a path to save model data. ''' return self.confs['data_dir'] @data_dir.setter def data_dir(self,values): self.confs['data_dir'] = values @property def source_volume(self): ''' The extent of source volume, in the form of ``[xmin,xmax,ymin,ymax,zmin,zmax]``. ''' return self.confs['source_volume'] @source_volume.setter def source_volume(self,value): self._source_volume = value self.confs['source_volume'] = value self.gen_mesh() @property def subtract_mean(self): return self.confs['subtract_mean'] @subtract_mean.setter def subtract_mean(self,values): self.confs['subtract_mean'] = values @property def weights(self): ''' inverse variance of each distributions. ''' return self._weights @weights.setter def weights(self,values): self._weights = values self.confs['weights'] = values if not self.kernel_op is None: self.kernel_op.weights = self._weights @property def optimize_keys(self): ''' inverse variance of each distributions. ''' return self.confs['optimize_keys'] @optimize_keys.setter def optimize_keys(self,values): self.confs['optimize_keys'] = values @property def smooth_components(self): ''' partial derivatives used as smooth components. Example: ``'dxx'`` means :math:`\frac{\partial^2 m}{\partial x^2}` ''' return self._smooth_components @smooth_components.setter def smooth_components(self,values): self._smooth_components = values self.confs['smooth_components'] = _smooth_components @property def refer_densities(self): '''reference density. The length of this vector should match the length of model density. ''' tmp = [] for density in self.constraints_val['refers']: tmp.append(density.reshape(self._nz,self._ny,self._nx)) return tmp @refer_densities.setter def refer_densities(self,value): tmp = [] for density in value: tmp.append(density.ravel()) self.constraints_val['refers'] = tmp @property def nzyx(self): '''model dimension, with the form of ``[nz,ny,nx]`` ''' return self.confs['nzyx'] @nzyx.setter def nzyx(self,values): self.confs['nzyx'] = values self._nz,self._ny,self._nx = values self.nobsx = values[2] self.nobsy = values[1] self.gen_model_name() if not self.constraints['depth'] is None: self.constraints['depth'] = self.constraints['depth'].reshape(self._nz,-1)[:,0]*np.ones((1,self._nx*self._ny)) self.constraints['depth'] = self.constraints['depth'].ravel() self.constraints['refers'] = np.ones(self._nx*self._ny*self._nz) @property def nx(self): ''' Number of cells along x-axis. ''' return self._nx @nx.setter def nx(self,value): self._nx = value self.nobsx = self._nx self.gen_model_name() if not self.constraints['depth'] is None: self.constraints['depth'] = self.constraints['depth'].reshape(self._nz,-1)[:,0]*np.ones((1,self._nx*self._ny)) self.constraints['depth'] = self.constraints['depth'].ravel() self.constraints['refers'] = np.ones(self._nx*self._ny*self._nz) @property def ny(self): ''' Number of cells along y-axis. ''' return self._ny @ny.setter def ny(self,value): self._ny = value self.nobsy = self._ny self.gen_model_name() if not self.constraints['depth'] is None: self.constraints['depth'] = self.constraints['depth'].reshape(self._nz,-1)[:,0]*np.ones((1,self._nx*self._ny)) self.constraints['depth'] = self.constraints['depth'].ravel() self.constraints['refers'] = np.ones(self._nx*self._ny*self._nz) @property def nz(self): ''' Number of cells along z-axis. ''' return self._nz @nz.setter def nz(self,value): self._nz = value self.gen_model_name() self.constraints['refers'] = np.ones(self._nx*self._ny*self._nz) print("Warning: nz changed. \nDon't forget setting depth constraints.") @property def model_density(self): ''' This vector is used for calculating gravity field, i.e. forward calculating. ''' return(self._model_density.reshape(self.nz,self.ny,self.nx)) @model_density.setter def model_density(self,value): self._model_density = value.ravel() self.confs['model_density'] = self._model_density def gen_mesh(self,height = -1): ''' Generate mesh of the model. Args: height (float): height of the observations. ''' shape = (self._nz, self._ny, self._nx) self.mesh = PrismMesh(self._source_volume, shape) density = np.ones(shape)*1.0e3 self.mesh.addprop('density', density.ravel()) # generate obs grid # coordinate: x North-South,y East-West # gridder is in the order: (nx,ny) self.obs_area = (self._source_volume[0]+0.5*self.mesh.dims[0], self._source_volume[1]-0.5*self.mesh.dims[0], self._source_volume[2]+0.5*self.mesh.dims[1], self._source_volume[3]-0.5*self.mesh.dims[1]) obs_shape = (self.nobsx, self.nobsy) self.xp, self.yp, self.zp = gridder.regular(self.obs_area, obs_shape, z=height) def _gen_walsh_matrix(self): '''generate walsh matrix in the order of sequence2. ''' print('generating walsh_matrix') if os.path.exists(self.fname): with h5py.File(self.fname,mode='r') as f: if not 'walsh_matrix' in f.keys(): have_walsh_matrix = False else: have_walsh_matrix = True else: have_walsh_matrix = False if have_walsh_matrix: return walsh_matrix = walsh.walsh_matrix(self._nx*self._ny*self._nz, normalized=True, ordering='sequence2', nxyz=(self._nx,self._ny,self._nz)) walsh_matrix = walsh_matrix.astype(np.float64) step = self._nx*self._ny*self._nz//4 components = ['0','1','2','3'] with h5py.File(self.fname,mode='a') as f: fgroup = f.create_group('walsh_matrix') for i in range(4): fgroup.create_dataset(components[i],data=walsh_matrix[i*step:(i+1)*step,:]) def gen_kernel(self, process = 1): '''generate kernel matrix. Because it is multilevel toeplitz matrix, only the first row are needed (i.e. all source cell contribution to the first observation position). .. note:: The coordinate system of the input parameters is to be x -> North, y -> East and z -> **DOWN**. .. note:: All input values in **SI** units(!) and output in **mGal**! ''' def calc_kernel(i): return prism.gz(self.xp[0:1],self.yp[0:1],self.zp[0:1],[self.mesh[i]]) if process > 1: #Winodws MP running has possiblely errors. print('Number of process:',process) with Pool(processes=process) as pool: kernel0 = pool.map(calc_kernel,range(len(self.mesh))) else: kernel0 = [calc_kernel(i) for i in range(len(self.mesh))] self.kernel0 = np.array(kernel0).reshape(self.nz,self.ny,self.nx) self.kernel_op = AbicLSQOperator(self.kernel0, depth_constraint=self.constraints['depth'], smooth_components=self._smooth_components, refer_constraints=self.constraints['refers'], weights=self._weights) def load_kernel(self,fname): '''load kernel matrix from file. Only the first row are needed. File format follows numpy's savetxt. ''' try: self.kernel0 = np.loadtxt(fname) except OSError: fname = pathlib.Path(self.data_dir)/pathlib.Path(fname) self.kernel0 = np.loadtxt(fname) self.kernel0 = np.array(self.kernel0).reshape(self.nz,self.ny,self.nx) self.kernel_op = AbicLSQOperator(self.kernel0, depth_constraint=self.constraints['depth'], smooth_components=self._smooth_components, refer_constraints=self.constraints['refers'], weights=self._weights) def _diagvec(self,vec=None,diag=None): if vec.ndim == 1: return vec * diag else: return vec * diag.reshape(1,-1) @timeit def walsh_transform(self,keys=None): '''walsh transform of kernel matrix and constraint matrices. ''' if keys is None: keys = ['kernel'] + list(self.constraints.keys()) + list(self._smooth_components) else: keys = keys if use_gpu > 0: import cupy as cp is_stored = dict() for key in keys: is_stored[key] = False if os.path.exists(self.fname): with h5py.File(self.fname,mode='r') as f: for key in keys: try: if '3' in f[key].keys(): is_stored[key] = True if key == 'depth': res = f['depth']['constraint'][:] - self.constraints['depth'] res = np.linalg.norm(res)/np.linalg.norm(self.constraints['depth']) if res > 1.0e-3: is_stored[key] = False if key == 'kernel': res = f['kernel']['source_volume'][:] - np.array(self.source_volume) res = np.linalg.norm(res)/np.linalg.norm(np.array(self.source_volume)) if res > 1.0e-3: is_stored[key] = False except KeyError: continue self._gen_walsh_matrix() logn = int(np.ceil(np.log2(self._nx*self._ny*self._nz))) norm_walsh = 1./(np.sqrt(2)**logn) blocks = ['0','1','2','3'] matvec_op = {'kernel':self.kernel_op.gtoep.matvec, 'depth': lambda x: self._diagvec(x,diag=np.sqrt(self.constraints['depth'])) } for key in self._smooth_components: matvec_op[key] = lambda x: self.smop.derivation(x.reshape(-1,self.nz,self.ny,self.nx),component=key).reshape(x.shape[0],-1) is_stored['refers'] = True for key in keys: if is_stored[key]: print('walsh transformation of {} already exists.'.format(key)) continue print('performing walsh transformation on {}.'.format(key)) step = self.nx*self.ny*self.nz // 4 if key == 'depth': step = self._nz with h5py.File(self.fname,mode='a') as f: try: del f[key] except KeyError: pass dxyz_group = f.create_group(key) walsh_group = f['walsh_matrix'] for i in range(4): print("\t progress {}/4".format(i)) part_walsh = walsh_group[blocks[i]][:] if key == 'depth': part_walsh = walsh_group[blocks[i]][:self._nz] part_walsh = matvec_op[key](part_walsh) if use_gpu > 0: with cp.cuda.Device(self.gpu_id): res = cp.zeros((step,step)) j = 0 while j*step < part_walsh.shape[1]: tmp_block_gpu = cp.asarray(part_walsh[:,j*step:(j+1)*step]) res += tmp_block_gpu @ tmp_block_gpu.T j += 1 res = cp.asnumpy(res) if key in self._smooth_components: res[np.abs(res)<1.0e-1*norm_walsh] = 0. tmp_block_gpu = None mempool = cp.get_default_memory_pool() pinned_mempool = cp.get_default_pinned_memory_pool() mempool.free_all_blocks() pinned_mempool.free_all_blocks() else: res = np.zeros((step,step)) j = 0 while j*step < part_walsh.shape[1]: tmp_block_gpu = np.asarray(part_walsh[:,j*step:(j+1)*step]) res += tmp_block_gpu @ tmp_block_gpu.T j += 1 if key in self._smooth_components: res[np.abs(res)<1.0e-1*norm_walsh] = 0. dxyz_group.create_dataset(blocks[i],data=res) if ('depth' in keys) and (not is_stored['depth']): with h5py.File(self.fname,mode='a') as f: try: del f['depth']['constraint'] except KeyError: pass dxyz_group = f['depth'] dxyz_group.create_dataset('constraint',data=self.constraints['depth']) if ('kernel' in keys) and (not is_stored['kernel']): with h5py.File(self.fname,mode='a') as f: try: del f['kernel']['source_volume'] except KeyError: pass dxyz_group = f['kernel'] dxyz_group.create_dataset('source_volume',data=np.array(self._source_volume)) @property def depth_constraint(self): '''Diagonal of the depth constraint matrix. One real number for each layer. Stored in an vector. ''' return(self.constraints['depth'].reshape(self._nz,-1)[:,0]) @depth_constraint.setter def depth_constraint(self,value): self.constraints['depth'] = (value.reshape(-1,1)*np.ones((1,self._nx*self._ny))).ravel() @timeit def forward(self,model_density=None): ''' Calculate gravity field from model_density. Args: model_density (np.Array): densities of each model cell. Reshaped from (nz,ny,nx) to (nz*ny*nx) ''' if model_density is None: model_density = self._model_density else: model_density = model_density.ravel() self.obs_data = self.kernel_op.gtoep.matvec(model_density) def _gen_rhs(self): '''generate right hand side of the least square equation of min_u. ''' self.rhs = self._weights['obs']*self.kernel_op.gtoep.rmatvec(self.obs_data) if 'refers' in self._weights.keys(): for i,refer_weight in enumerate(self._weights['refers']): v = self.constraints_val['refers'][i].ravel() if 'depth' in self._weights.keys(): v = self.constraints['depth']*v self.rhs += (refer_weight *self.constraints['depth'] *v) if self.smooth_on == 'm-m0': for key in self._smooth_components: tmp2 = v.reshape(-1,self._nz,self._ny,self._nx) tmp2 = self.smop.derivation(tmp2,component=key) tmp2 = self.smop.rderivation(tmp2,component=key) if v.ndim == 1: self.rhs += self._weights[key]*self._weights['depth']*self.constraints['depth']*tmp2.ravel() else: self.rhs += self._weights[key]*self._weights['depth']*self.constraints['depth']*tmp2.reshape(v.shape[0],-1) @timeit def do_linear_solve(self,tol=1.0e-5): ''' solve the least square equation of min_u. Args: tol (float): tol of CG algorithm. ''' self.do_linear_solve_quiet(tol=tol) def do_linear_solve_quiet(self,tol=1.0e-5): ''' solve the least square equation of min_u with minimum of message printed. Args: tol (float): tol of CG algorithm. ''' self._gen_rhs() if self.subtract_mean: sum_obs = np.sum(self.obs_data) tmp_b = np.zeros(len(self.rhs)+1) #np.zeros(len(self.rhs+1)) #chenshi tmp_b[:-1] = self.rhs tmp_b[-1] = sum_obs tmp_op = AbicLSQOperator2(self.kernel_op) self.solution = spsparse.linalg.cg(tmp_op,tmp_b,tol=tol)[0] else: self.solution = spsparse.linalg.cg(self.kernel_op,self.rhs,tol=tol)[0] @timeit def calc_u(self,solved=False,x=None): '''calc min_u value. Args: solved (bool): whether or not the least square equation solved already. x (array): use this x to calculate value of U instead of doing optimization. ''' return self.calc_u_quiet(solved,x) @timeit def calc_min_u(self,solved=False,x=None): '''Another name of calc_u. We keep this function for backward compatability. ''' return self.calc_u_quiet(solved,x) def calc_u_quiet(self,solved=False,x=None): '''calc min_u value with minimum message printed out. Args: solved (bool): whether or not the least square equation solved already. x (array): use this x to calculate value of U instead of doing optimization. ''' if x is None: if not solved: self.do_linear_solve_quiet() x = self.solution self.min_u_val = self._weights['obs']*np.linalg.norm(self.kernel_op.gtoep.matvec(x) - self.obs_data)**2 if ('refers' in self._weights.keys()) and (self.smooth_on == 'm-m0'): v = x - self.constraints_val['refer'] else: v = x if 'depth' in self._weights.keys(): v = self.constraints['depth']*v for key in self._smooth_components: tmp2 = self.smop.derivation(v.reshape(self._nz,self._ny,self._nx), component=key) self.min_u_val += self._weights[key]*np.linalg.norm(tmp2.ravel())**2 if 'refers' in self._weights.keys(): for i,refer_density in enumerate(self.constraints_val['refers']): v = x - refer_density if 'depth' in self._weights.keys(): v = self.constraints['depth']*v self.min_u_val += self._weights['refers'][i] *np.linalg.norm(v)**2 return self.min_u_val def jac_u(self,x=None): ''' jacobian of the function U. ''' res = self.kernel_op.matvec(x) - self.rhs return 2.*res def hessp_u(self,x,v): '''Hessian of the function U. ''' res = self.kernel_op.matvec(v) return 2.*res @timeit def bound_optimize(self,x0=None): '''optimize function U using boundary constraint. ''' density_bounds = Bounds(self.min_density,self.max_density) if x0 is None: x0 = np.zeros(self._nx*self._ny*self._nz)+(self.max_density - self.min_density)/2. self.bound_solution = minimize(lambda x:self.calc_u_quiet(solved=True,x=x), x0, method='trust-constr', jac=self.jac_u, hessp=self.hessp_u, bounds=density_bounds, ) def lasso_target(self,x): '''Minimize the 1-norm instead 2-norm of the model equation. This function is used to the target function. ''' # self.min_u_val = self._weights['obs']*np.linalg.norm(self.kernel_op.gtoep.matvec(x) - self.obs_data)**2 self.min_u_val = self._weights['obs']*np.linalg.norm(self.kernel_op.gtoep.matvec(x) - self.obs_data) if ('refer' in self._weights.keys()) and (self.smooth_on == 'm-m0'): v = x - self.constraints_val['refer'] else: v = x if 'depth' in self._weights.keys(): v = np.sqrt(self._weights['depth'])*self.constraints['depth']*v for key in self._smooth_components: tmp2 = self.smop.derivation(v.reshape(self._nz,self._ny,self._nx), component=key) self.min_u_val += self._weights[key]*np.linalg.norm(tmp2.ravel()) if 'refer' in self._weights.keys(): v = x - self.constraints_val['refer'] if 'depth' in self._weights.keys(): v = np.sqrt(self._weights['depth'])*self.constraints['depth']*v self.min_u_val += self._weights['refer'] *np.linalg.norm(v) return self.min_u_val def lasso_jac(self,x): '''jacobian of the lasso target function. ''' # jac = self.kernel_op.gtoep.rmatvec(self.kernel_op.gtoep.matvec(x)) - self.kernel_op.gtoep.rmatvec(self.obs_data) # jac = 2.0*self._weights['obs']*jac jac = self.kernel_op.gtoep.rmatvec(self.kernel_op.gtoep.matvec(x)) - self.kernel_op.gtoep.rmatvec(self.obs_data) jac = self._weights['obs']*jac norm_res = np.linalg.norm(self.obs_data - self.kernel_op.gtoep.matvec(x)) jac = jac/norm_res if 'refer' in self.weights.keys(): v = x - self.constraints_val['refer'] if 'depth' in self._weights.keys(): v = self.constraints['depth']*v norm_refer = np.linalg.norm(v) jac += self.weights['refer']*self.weights['depth']*self.constraints['depth']*v/norm_refer for key in self.smooth_components: v = x if 'depth' in self._weights.keys(): v = self.constraints['depth']*v tmp2 = self.smop.derivation(v.reshape(self._nz,self._ny,self._nx), component=key) smooth_norm = np.linalg.norm(tmp2.ravel()) tmp2 = self.smop.rderivation(tmp2,component=key) jac += self.weights[key]*self.weights['depth']*self.constraints['depth']*tmp2.ravel()/smooth_norm return jac def lasso_hessp(self,x,v): '''hessian of the lasso target function. ''' # res = self.kernel_op.gtoep.rmatvec(self.kernel_op.gtoep.matvec(v)) # res = 2.0*self._weights['obs']*res norm_res = np.linalg.norm(self.obs_data - self.kernel_op.gtoep.matvec(x)) res = self.kernel_op.gtoep.rmatvec(self.kernel_op.gtoep.matvec(v))/norm_res gradient_res = (self.kernel_op.gtoep.rmatvec(self.kernel_op.gtoep.matvec(x)) - self.kernel_op.gtoep.rmatvec(self.obs_data)) res -= np.dot(gradient_res,v)/norm_res**3 * gradient_res res *= self._weights['obs'] if 'refer' in self.weights.keys(): v2 = x - self.constraints_val['refer'] if 'depth' in self._weights.keys(): v2 = self.constraints['depth']*v2 norm_refer = np.linalg.norm(v2) res += self.weights['refer']*self.weights['depth']/norm_refer*self.constraints['depth']*v grad_ref = self.constraints['depth']*v2 res -= (self.weights['refer']*self.weights['depth'] *(np.dot(v,grad_ref)/norm_refer**3*grad_ref)) for key in self.smooth_components: v2 = x if 'depth' in self._weights.keys(): v2 = self.constraints['depth']*v2 tmp2 = self.smop.derivation(v2.reshape(self._nz,self._ny,self._nx), component=key) smooth_norm = np.linalg.norm(tmp2.ravel()) tmp2 = self.smop.rderivation(tmp2,component=key) grad_sm = self.constraints['depth']*tmp2.ravel() tmp2 = v.reshape(self.nz,self.ny,self.nx) tmp2 = self.smop.derivation(tmp2,component=key) tmp2 = self.smop.rderivation(tmp2,component=key) tmp2 = self.constraints['depth']*tmp2.ravel() res += (self._weights['depth']*self._weights[key]/smooth_norm *(tmp2 - np.dot(v,grad_sm)*grad_sm/smooth_norm**2)) return res @timeit def lasso_optimize(self,x0=None): '''optimize the lasso function. ''' density_bounds = Bounds(self.min_density,self.max_density) if x0 is None: x0 = (np.random.rand(self._nx*self._ny*self._nz) *(self.max_density - self.min_density)/2. +(self.max_density + self.min_density)/2.) self.bound_solution = minimize(self.lasso_target, x0, method='trust-constr', jac=self.lasso_jac, hessp=self.lasso_hessp, bounds=density_bounds ) def calc_res(self): '''calculate the residual information including residuals and stds. The result of this function are stored in ``self.residuals`` and ``self.stds``. ''' self.residuals = dict() self.stds = dict() self.residuals['obs'] = np.linalg.norm(self.kernel_op.gtoep.matvec(self.solution)-self.obs_data)**2 self.stds['obs'] = np.std(self.kernel_op.gtoep.matvec(self.solution)-self.obs_data) for key in self._smooth_components: tmp2 = self.solution.reshape(self._nz,self._ny,self._nx) if ('refers' in self.constraints_val.keys()) and (self.smooth_on == 'm-m0'): tmp2 -= self.constraints_val['refer'].reshape(self._nz,self._ny,self._nx) tmp2 = self.smop.derivation(tmp2,component=key) self.residuals[key] = np.linalg.norm(tmp2.ravel())**2 self.stds[key] = np.std(tmp2.ravel()) if 'refers' in self.constraints_val.keys(): self.residuals['refers'] = [] self.stds['refers'] = [] for i,refer_density in self.constraints_val['refers']: self.residuals['refers'].append(np.linalg.norm(self.solution.ravel()-refer_density.ravel())**2) self.stds['refers'].append(np.std(self.solution.ravel()-refer_density.ravel())) def calc_log_prior_total_det_naive(self): '''calculate the determinant of prior distribution and joint distribution with minimum message printed out. ''' self.log_prior_det_val = 0 self.log_total_det_val = 0 prior_eigs = np.zeros(self._nx*self._ny*self._nz) total_eigs = np.zeros(self._nx*self._ny*self._nz) tmp_mat = np.zeros((self._nz*self._ny*self._nx,self._nz*self._ny*self._nx)) for dxyz_name in self._smooth_components: tmp_mat += self._weights[dxyz_name]*self.matrices[dxyz_name] if 'depth' in self._weights.keys(): tmp_mat = self.constraints['depth'].reshape(-1,1)*self.constraints['depth'].reshape(1,-1)*tmp_mat prior_eigs = np.linalg.svd(tmp_mat,compute_uv=False) eps = prior_eigs.max() * len(prior_eigs) * np.finfo(np.float64).eps eigs = prior_eigs[prior_eigs>eps] self.log_prior_det_val = np.sum(np.log(eigs)) tmp_mat += self._weights['obs']*self.matrices['obs'] if 'refers' in self._weights.keys(): tmp_mat += sum(self._weights['refers'])*np.diag(self.constraints['depth'])**2 uu,total_eigs,vv = np.linalg.svd(tmp_mat,compute_uv=True) self.log_total_det_val = np.sum(np.log(total_eigs)) self._gen_rhs() self.solution = np.zeros(np.prod(self.nzyx)) self.solution = (vv.T @ ((1./total_eigs).ravel() * (uu.T @ self.rhs))) self.eigs = {'prior':prior_eigs,'total':total_eigs} return self.log_prior_det_val,self.log_total_det_val def calc_log_prior_total_det_walsh(self): '''calculate the determinant of prior distribution and joint distribution with minimum message printed out. ''' self.log_prior_det_val = 0 self.log_total_det_val = 0 blocks = ['0','1','2','3'] prior_eigs = np.zeros(self._nx*self._ny*self._nz) total_eigs = np.zeros(self._nx*self._ny*self._nz) step = self._nx*self._ny*self._nz//4 try: depth_weight = self._weights['depth'] except KeyError: depth_weight = 1. with h5py.File(self.fname,mode='r') as f: if 'depth' in self._weights.keys(): depth_walsh = f['depth']['0'][:] self._gen_rhs() self.solution = np.zeros(np.prod(self.nzyx)) for i_b,block in enumerate(blocks): tmp_block = np.zeros((step,step)) walsh_group = f['walsh_matrix'] part_walsh = walsh_group[blocks[i_b]][:] for dxyz_name in self._smooth_components: try: dxyz_walsh = f[dxyz_name][block][:].reshape(step//self._nz, self._nz, step//self._nz, self._nz) ein_path = np.einsum_path('mi,xiyj,jn->xmyn', depth_walsh.T, dxyz_walsh, depth_walsh, optimize='optimal')[0] tmp_multi = np.einsum('mi,xiyj,jn->xmyn', depth_walsh.T, dxyz_walsh, depth_walsh, optimize=ein_path) tmp_block += depth_weight*self._weights[dxyz_name]*tmp_multi.reshape(step,step) except KeyError: pass if use_gpu > 0: import cupy as cp with cp.cuda.Device(self.gpu_id): tmp_block_gpu = cp.asarray(tmp_block,dtype=np.float64) eigs = cp.linalg.svd(tmp_block_gpu,compute_uv=False) prior_eigs[i_b*step:(i_b+1)*step] = cp.asnumpy(eigs) eps = eigs.max() * len(eigs) * np.finfo(np.float64).eps eigs = eigs[eigs>eps] self.log_prior_det_val += cp.asnumpy(cp.sum(cp.log(eigs))) tmp_block_gpu = None eigs = None free_gpu() tmp_block += self._weights['obs']*f['kernel'][block][:] if 'refers' in self._weights.keys(): tmp_multi_small = depth_walsh.T@depth_walsh for i in range(step//self._nz): tmp_block[i*self._nz:(i+1)*self._nz, i*self._nz:(i+1)*self._nz] += sum(self._weights['refers'])*tmp_multi_small with cp.cuda.Device(self.gpu_id): tmp_block_gpu = cp.asarray(tmp_block,dtype=np.float64) eigs = cp.linalg.svd(tmp_block_gpu,compute_uv=False) total_eigs[i_b*step:(i_b+1)*step] = cp.asnumpy(eigs) #eigs = eigs[eigs>1.0e-12] self.log_total_det_val += cp.asnumpy(cp.sum(cp.log(eigs))) tmp_block_gpu = None eigs = None free_gpu() else: tmp_block_gpu = np.asarray(tmp_block,dtype=np.float64) eigs = np.linalg.svd(tmp_block_gpu,compute_uv=False) prior_eigs[i_b*step:(i_b+1)*step] = eigs eps = eigs.max() * len(eigs) * np.finfo(np.float64).eps eigs = eigs[eigs>eps] self.log_prior_det_val += np.sum(np.log(eigs)) tmp_block_gpu = None eigs = None tmp_block += self._weights['obs']*f['kernel'][block][:] if 'refers' in self._weights.keys(): tmp_multi_small = depth_walsh.T@depth_walsh #tmp_multi_small = np.eye(tmp_multi_small.shape[0]) for i in range(step//self._nz): tmp_block[i*self._nz:(i+1)*self._nz, i*self._nz:(i+1)*self._nz] += sum(self._weights['refers'])*tmp_multi_small tmp_block_gpu = np.asarray(tmp_block,dtype=np.float64) uu,eigs,vv = np.linalg.svd(tmp_block_gpu,compute_uv=True) total_eigs[i_b*step:(i_b+1)*step] = eigs #eigs = eigs[eigs>1.0e-12] self.log_total_det_val += np.sum(np.log(eigs)) tmp_block_gpu = None self.solution += part_walsh.T @ (vv.T @ ((1./eigs).ravel() * (uu.T @ (part_walsh @ self.rhs)))) eigs = None if use_gpu > 0: self.log_prior_det_val = cp.asnumpy(self.log_prior_det_val) self.log_total_det_val = cp.asnumpy(self.log_total_det_val) else: self.log_prior_det_val = self.log_prior_det_val self.log_total_det_val = self.log_total_det_val self.eigs = {'prior':prior_eigs,'total':total_eigs} return self.log_prior_det_val,self.log_total_det_val def calc_log_prior_total_det_quiet(self,mode=None): '''calculate the determinant of prior distribution and joint distribution with minimum message printed out. ''' if mode is None: mode = self.mode if mode == 'walsh': self.calc_log_prior_total_det_walsh() elif mode == 'naive': self.calc_log_prior_total_det_naive() else: raise ValueError('mode={} is not implemented!!'.format(mode)) def prepare_det(self,mode=None): if mode is None: mode = self.mode if mode == 'walsh': self.walsh_transform() elif mode == 'naive': self.matrices = dict() n_cell = self._nx*self._ny*self._nz self.matrices['G'] = self.kernel_op.gtoep.matvec(np.eye(self._nx*self._ny*self._nz)).T self.matrices['obs'] = self.kernel_op.gtoep.rmatvec(self.kernel_op.gtoep.matvec(np.eye(n_cell))) for key in self._smooth_components: self.matrices[key] = self.smop.rderivation(self.smop.derivation(np.eye(n_cell).reshape(-1,self._nz,self._ny,self._nx),key),key).reshape(n_cell,n_cell) else: raise ValueError('mode={} is not implemented!!'.format(mode)) @timeit def calc_log_prior_total_det(self,mode=None): '''calculate the determinant of prior distribution and joint distribution. ''' return self.calc_log_prior_total_det_quiet(mode=mode) def calc_log_obs_det_quiet(self): '''calculate the determinant of observation's distribution with minimum message printed out. ''' self.log_obs_det_val = (4*sum(np.log(self.constraints['depth'])) +np.prod(self.nzyx)*(np.log(self.weights['refers'][0]) +np.log(self.weights['refers'][1])) +np.log(self._weights['obs'])*len(self.obs_data)) return self.log_obs_det_val @timeit def calc_log_obs_det(self): '''calculate the determinant of observation's distribution message printed out. ''' return self.calc_log_obs_det_quiet() @timeit def calc_abic(self,mode=None): '''calculate abic value: -log_prior_det_value+log_total_det-log_obs_det+min_u''' self.calc_log_obs_det() self.calc_log_prior_total_det(mode=mode) self.calc_u(solved=True) self.abic_val = (self.log_total_det_val + self.min_u_val - self.log_prior_det_val - self.log_obs_det_val) return self.abic_val def calc_abic_quiet(self,mode=None): '''calculate abic value: -log_prior_det_value+log_total_det-log_obs_det+min_u''' self.calc_log_obs_det_quiet() self.calc_log_prior_total_det_quiet(mode=mode) self.calc_u_quiet(solved=True) self.abic_val = (self.log_total_det_val + self.min_u_val - self.log_prior_det_val - self.log_obs_det_val) self.optimize_log['parameters'].append(self.weights) self.optimize_log['abic_vals'].append(self.abic_val) return self.abic_val def _abic_optimize_exp(self,mode=None,**opt_args): '''optimize the abic value. Use the log and exp trick to constrain weights to be positive. ''' #optimize_keys = list(set(self._weights.keys())-set(['depth'])) if self.confs['optimize_keys'] is None: optimize_keys = list(self._weights.keys()) else: optimize_keys = self.confs['optimize_keys'] optimize_keys_1 = list(set(optimize_keys)-set(['refers'])) def abic_target(x): tmp_weights = dict() for i,key in enumerate(optimize_keys_1): tmp_weights[key] = np.exp(x[i]) if 'refers' in optimize_keys: tmp_weights['refers'] = list(np.exp(x[-len(self.weights['refers']):])) self.weights = {**self.weights,**tmp_weights} self.calc_abic_quiet(mode=mode) return self.abic_val n_weights = len(optimize_keys) if 'refers' in optimize_keys: n_weights += len(self.weights['refers']) - 1 x0 = np.zeros(n_weights) for i,key in enumerate(optimize_keys_1): x0[i] = np.log(self._weights[key]) if 'refers' in optimize_keys: for i,refer_weight in enumerate(self._weights['refers']): x0[len(optimize_keys_1)+i] = np.log(refer_weight) self.abic_optimize_summary = minimize(abic_target, x0, **opt_args) def _abic_optimize_bound(self): '''optimize the abic value. Use scipy's COBYLA algorithm to constrain weights to be positive. ''' if self.confs['optimize_keys'] is None: optimize_keys = list(self._weights.keys()) else: optimize_keys = self.confs['optimize_keys'] def abic_target(x): tmp_weights = dict() for i,key in enumerate(optimize_keys): tmp_weights[key] = np.exp(x[i]) self.weights = {**self.weights,**tmp_weights} return self.calc_abic_quiet() x0 = np.zeros(len(optimize_keys)) for i,key in enumerate(optimize_keys): x0[i] = self._weights[key] weight_constraint = LinearConstraint(np.eye(len(optimize_keys)),0.,np.inf) self.abic_optimize_summary = minimize(abic_target, x0, method='COBYLA', constraints=weight_constraint) @timeit def abic_optimize(self,mode=None,**opt_args): '''optimize the abic value. This is the interface for users. ''' self._abic_optimize_exp(mode=mode,**opt_args) @timeit def para_grad(self,x): pass def u_bound(self): pass def print_summary(self): print('abic values:{}'.format(self.abic_val)) print('log total det:{}'.format(self.log_total_det_val)) print('log prior det:{}'.format(self.log_prior_det_val)) print('log obs det:{}'.format(self.log_obs_det_val)) print('min u:{}'.format(self.min_u_val)) print('std:',end=' ') print(self.stds) print('1/var:',end=' ') print({k:1./v**2 for k,v in self.stds.items()}) print('norms:',end=' ') print(self.residuals)
class GravInvAbicModel: def __init__(self, nzyx=[4,4,4], smooth_components=['dx','dy','dz'], depth_constraint=None, model_density=None, refer_density=None, weights=None, source_volume=None, smooth_on='m', data_dir='/data/gravity_inversion'): self._nz,self._ny,self._nx = nzyx self.smooth_on = smooth_on self.dxyz_shapes = {'dx':(self._nz,self._ny,self._nx), 'dy':(self._nz,self._nx*self._ny), 'dz':(self._nx*self._ny*self._nz,)} self.dxyz_spaces = {'dx':self._nx-1, 'dy':self._nx*(self._ny-1), 'dz':self._nx*self._ny*(self._nz-1)} self.data_dir = data_dir self.gen_model_name() self.nobsx = nzyx[2] self.nobsy = nzyx[1] self.source_volume = source_volume if model_density is None: self._model_density = None else: self._model_density = model_density.ravel() self._smooth_components = smooth_components self.constraints = dict() self.constraints_val = dict() if depth_constraint is None: self.constraints['depth'] = np.ones(np.prod(nzyx)) self.constraints_val['depth'] = None else: self.constraints['depth'] = (depth_constraint.reshape(-1,1)*np.ones((1,self._nx*self._ny))).ravel() self.constraints_val['depth'] = 0 if refer_density is None: self.constraints['refer'] = None self.constraints_val['refer'] = None else: self.constraints['refer'] = np.ones(self._nx*self._ny*self._nz) self.constraints_val['refer'] = refer_density.ravel() self._weights = weights if not 'depth' in self._weights.keys(): self._weights['depth'] = 1.0 self._gen_dxyz_constraint() self.kernel_op = None self.abic_val = 0 self.log_total_det_val = 0 self.log_prior_det_val = 0 self.log_obs_det_val = 0 self.min_u_val = 0 self.min_density = -1.0e4 self.max_density = 1.0e4 @property def source_volume(self): return self._source_volume @source_volume.setter def source_volume(self,value): self._source_volume = value self.gen_mesh() def gen_model_name(self): self.model_name = '{}x{}x{}'.format(self._nx,self._ny,self._nz) self.fname = pathlib.Path(self.data_dir)/pathlib.Path(self.model_name+'.h5') @property def weights(self): return self._weights @weights.setter def weights(self,values): self._weights = values if not self.kernel_op is None: self.kernel_op.weights = self._weights @property def smooth_components(self): return self._smooth_components @smooth_components.setter def smooth_components(self,values): self._smooth_components = values self._gen_dxyz_constraint() if not self.kernel_op is None: self.kernel_op.dxyz_constraint = self.dxyz_constraint @timeit def _gen_dxyz_constraint(self): '''first generate multi-level circulant matrix, constraint of dx is a part of it. then calculate it's eigenvalues. self._dx stores the eigenvalues finally. When multiply it with a vector, specific element should be discarded''' self.dxyz_constraint = dict() for component in self._smooth_components: tmp = np.zeros(self.nx*self.ny*self.nz) tmp[0] = 1 tmp[self.dxyz_spaces[component]] = -1 tmp = tmp.reshape(self.dxyz_shapes[component]) self.dxyz_constraint[component] = np.fft.fftn(tmp) self.constraints[component] = self.dxyz_constraint[component] @property def refer_density(self): return self.constraints_val['refer'].reshape(self._nz,self._ny,self._nx) @refer_density.setter def refer_density(self,value): self.constraints_val['refer'] = value.ravel() @property def nx(self): return self._nx @nx.setter def nx(self,value): self._nx = value self.nobsx = self._nx self.gen_model_name() if not self.constraints['depth'] is None: self.constraints['depth'] = self.constraints['depth'].reshape(self._nz,-1)[:,0]*np.ones((1,self._nx*self._ny)) self.constraints['depth'] = self.constraints['depth'].ravel() self.constraints['refer'] = np.ones(self._nx*self._ny*self._nz) @property def ny(self): return self._ny @ny.setter def ny(self,value): self._ny = value self.nobsy = self._ny self.gen_model_name() if not self.constraints['depth'] is None: self.constraints['depth'] = self.constraints['depth'].reshape(self._nz,-1)[:,0]*np.ones((1,self._nx*self._ny)) self.constraints['depth'] = self.constraints['depth'].ravel() self.constraints['refer'] = np.ones(self._nx*self._ny*self._nz) @property def nz(self): return self._nz @nz.setter def nz(self,value): self._nz = value self.gen_model_name() self.constraints['refer'] = np.ones(self._nx*self._ny*self._nz) print("Warning: nz changed. \nDon't forget setting depth constraints.") @property def model_density(self): return(self._model_density.reshape(self.nz,self.ny,self.nx)) @model_density.setter def model_density(self,value): self._model_density = value.ravel() def gen_mesh(self,height = -1): shape = (self._nz, self._ny, self._nx) self.mesh = PrismMesh(self._source_volume, shape) density = np.ones(shape)*1.0e3 self.mesh.addprop('density', density.ravel()) # generate obs grid # coordinate: x North-South,y East-West # gridder is in the order: (nx,ny) self.obs_area = (self._source_volume[0]+0.5*self.mesh.dims[0], self._source_volume[1]-0.5*self.mesh.dims[0], self._source_volume[2]+0.5*self.mesh.dims[1], self._source_volume[3]-0.5*self.mesh.dims[1]) obs_shape = (self.nobsx, self.nobsy) self.xp, self.yp, self.zp = gridder.regular(self.obs_area, obs_shape, z=height) def _gen_walsh_matrix(self): print('generating walsh_matrix') if os.path.exists(self.fname): with h5py.File(self.fname,mode='r') as f: if not 'walsh_matrix' in f.keys(): have_walsh_matrix = False else: have_walsh_matrix = True else: have_walsh_matrix = False if have_walsh_matrix: return walsh_matrix = walsh.walsh_matrix(self._nx*self._ny*self._nz, normalized=True, ordering='sequence2', nxyz=(self._nx,self._ny,self._nz)) walsh_matrix = walsh_matrix.astype(np.float32) step = self._nx*self._ny*self._nz//4 components = ['0','1','2','3'] with h5py.File(self.fname,mode='a') as f: fgroup = f.create_group('walsh_matrix') for i in range(4): fgroup.create_dataset(components[i],data=walsh_matrix[i*step:(i+1)*step,:]) def gen_kernel(self): def calc_kernel(i): return prism.gz(self.xp[0:1],self.yp[0:1],self.zp[0:1],[self.mesh[i]]) with Pool(processes=16) as pool: kernel0 = pool.map(calc_kernel,range(len(self.mesh))) self.kernel0 = np.array(kernel0).reshape(self.nz,self.ny,self.nx) self.kernel_op = AbicLSQOperator(self.kernel0, depth_constraint=self.constraints['depth'], dxyz_constraint=self.dxyz_constraint, refer_constraint=self.constraints['refer'], weights=self._weights) def _dxyzvec(self,vec=None,key=None): res = vec.reshape(-1,*self.dxyz_shapes[key]) axes = np.arange(1,res.ndim) res = np.fft.ifftn(np.fft.fftn(res,axes=axes)*self.dxyz_constraint[key],axes=axes).real slices = [slice(None)]*res.ndim slices[-1] = slice(0,self.dxyz_spaces[key]) if vec.ndim == 1: return res[tuple(slices)].ravel() else: return res[tuple(slices)].reshape(vec.shape[0],-1) def _diagvec(self,vec=None,diag=None): if vec.ndim == 1: return vec * diag else: return vec * diag.reshape(1,-1) @timeit def walsh_transform(self,keys=None): if keys is None: keys = ['kernel'] + list(self.constraints.keys()) else: keys = keys is_stored = dict() for key in keys: is_stored[key] = False if os.path.exists(self.fname): with h5py.File(self.fname,mode='r') as f: for key in keys: try: if '3' in f[key].keys(): is_stored[key] = True if key == 'depth': res = f['depth']['constraint'][:] - self.constraints['depth'] res = np.linalg.norm(res)/np.linalg.norm(self.constraints['depth']) if res > 1.0e-3: is_stored[key] = False except KeyError: continue self._gen_walsh_matrix() logn = int(np.ceil(np.log2(self._nx*self._ny*self._nz))) norm_walsh = 1./(np.sqrt(2)**logn) blocks = ['0','1','2','3'] matvec_op = {'kernel':self.kernel_op.gtoep.matvec, 'dx': lambda x: self._dxyzvec(x,key='dx'), 'dy': lambda x: self._dxyzvec(x,key='dy'), 'dz': lambda x: self._dxyzvec(x,key='dz'), 'refer': lambda x: self._diagvec(x,diag=self.constraints['refer']), 'depth': lambda x: self._diagvec(x,diag=np.sqrt(self.constraints['depth'])) } is_stored['refer'] = True for key in keys: if is_stored[key]: print('walsh transformation of {} already exists.'.format(key)) continue print('performing walsh transformation on {}.'.format(key)) step = self.nx*self.ny*self.nz // 4 if key == 'depth': step = self._nz with h5py.File(self.fname,mode='a') as f: try: del f[key] except KeyError: pass dxyz_group = f.create_group(key) walsh_group = f['walsh_matrix'] for i in range(4): print("\t progress {}/4".format(i)) part_walsh = walsh_group[blocks[i]][:] if key == 'depth': part_walsh = walsh_group[blocks[i]][:self._nz] part_walsh = matvec_op[key](part_walsh) with cp.cuda.Device(2): res = cp.zeros((step,step)) j = 0 while j*step < part_walsh.shape[1]: tmp_block_gpu = cp.asarray(part_walsh[:,j*step:(j+1)*step]) res += tmp_block_gpu @ tmp_block_gpu.T j += 1 res = cp.asnumpy(res) if key in self._smooth_components: res[np.abs(res)<1.0e-1*norm_walsh] = 0. tmp_block_gpu = None mempool = cp.get_default_memory_pool() pinned_mempool = cp.get_default_pinned_memory_pool() mempool.free_all_blocks() pinned_mempool.free_all_blocks() dxyz_group.create_dataset(blocks[i],data=res) if ('depth' in keys) and (not is_stored['depth']): with h5py.File(self.fname,mode='a') as f: try: del f['depth_constraint'] except KeyError: pass dxyz_group = f['depth'] dxyz_group.create_dataset('constraint',data=self.constraints['depth']) @property def depth_constraint(self): return(self.constraints['depth'].reshape(self._nz,-1)[:,0]) @depth_constraint.setter def depth_constraint(self,value): self.constraints['depth'] = (value.reshape(-1,1)*np.ones((1,self._nx*self._ny))).ravel() @timeit def forward(self,model_density=None): if model_density is None: model_density = self._model_density else: model_density = model_density.ravel() self.obs_data = self.kernel_op.gtoep.matvec(model_density) def _gen_rhs(self): self.rhs = self._weights['obs']*self.kernel_op.gtoep.rmatvec(self.obs_data) if 'depth' in self._weights.keys(): v = self.constraints['depth']*self.constraints_val['refer'] if 'refer' in self._weights.keys(): self.rhs += (self._weights['refer'] *self._weights['depth'] *self.constraints['depth'] *v) if self.smooth_on == 'm-m0': if not self.dxyz_constraint is None: for key,constraint in self.dxyz_constraint.items(): if not key in self._weights.keys(): continue tmp2 = v.reshape(-1,*constraint.shape) fft_comp = list(range(tmp2.ndim))[1:] tmp2 = np.fft.ifftn(np.fft.fftn(tmp2,axes=fft_comp)*constraint,axes=fft_comp).real slices = [slice(None)]*tmp2.ndim slices[-1] = slice(self.dxyz_spaces[key],None) tmp2[tuple(slices)] = 0 tmp2 = np.real(np.fft.ifftn(np.fft.fftn(tmp2,axes=fft_comp)*np.conj(constraint),axes=fft_comp)) if v.ndim == 1: self.rhs += self._weights[key]*self._weights['depth']*self.constraints['depth']*tmp2.ravel() else: self.rhs += self._weights[key]*self._weights['depth']*self.constraints['depth']*tmp2.reshape(v.shape[0],-1) @timeit def do_linear_solve(self): self._gen_rhs() self.solution = spsparse.linalg.cg(self.kernel_op,self.rhs,tol=1.0e-5)[0] @timeit def calc_min_u(self,solved=False,x=None): if x is None: if not solved: self.do_linear_solve() x = self.solution self.min_u_val = self._weights['obs']*np.linalg.norm(self.kernel_op.gtoep.matvec(x) - self.obs_data)**2 if ('refer' in self._weights.keys()) and (self.smooth_on == 'm-m0'): v = x - self.constraints_val['refer'] else: v = x if 'depth' in self._weights.keys(): v = np.sqrt(self._weights['depth'])*self.constraints['depth']*v if not self.dxyz_constraint is None: for key,constraint in self.dxyz_constraint.items(): if not key in self._weights.keys(): continue tmp2 = np.fft.ifftn( np.fft.fftn(v.reshape(constraint.shape))*constraint ).real slices = [slice(None)]*constraint.ndim slices[-1] = slice(0,self.dxyz_spaces[key]) self.min_u_val += self._weights[key]*np.linalg.norm(tmp2[tuple(slices)].ravel())**2 if 'refer' in self._weights.keys(): v = x - self.constraints_val['refer'] if 'depth' in self._weights.keys(): v = np.sqrt(self._weights['depth'])*self.constraints['depth']*v self.min_u_val += self._weights['refer'] *np.linalg.norm(v)**2 return self.min_u_val def bound_constraint_u(self,x=None): self.calc_min_u(x=x,solved=True) log_barrier = np.sum(np.log(x-self.min_density) + np.log(self.max_density-x)) return self.min_u_val - 2.*self._weights['bound']*log_barrier def bound_jac_u(self,x=None): res = 0. res += self._weights['obs']*(self.kernel_op.gtoep.matvec(x) - self.obs_data) if ('refer' in self._weights.keys()) and (self.smooth_on == 'm-m0'): v = x - self.constraints_val['refer'] else: v = x if 'depth' in self._weights.keys(): v = self._weights['depth']*self.constraints['depth']*v if not self.dxyz_constraint is None: for key,constraint in self.dxyz_constraint.items(): if not key in self._weights.keys(): continue tmp2 = np.fft.ifftn( np.fft.fftn(v.reshape(constraint.shape))*constraint ).real slices = [slice(None)]*constraint.ndim slices[-1] = slice(0,self.dxyz_spaces[key]) res += self._weights[key]*tmp2[tuple(slices)].ravel() if 'refer' in self._weights.keys(): v = x - self.constraints_val['refer'] if 'depth' in self._weights.keys(): v = self._weights['depth']*self.constraints['depth']*v res += self._weights['refer'] *v res += self._weights['bound']*(1./(self.max_density-x) - 1./(x-self.min_density)) return 2.*res def bound_hessp_u(self,x,v): res = self.kernel_op.matvec(v) hess_diag = 1./(self.max_density-x)**2 + 1./(x-self.min_density)**2 res += self._weights['bound']*hess_diag*v return 2.*res def bound_optimize(self,x0=None): if x0 is None: if 'refer' in self._weights.keys(): x0 = self.constraints_val['refer'] else: x0 = np.zeros(self._nx*self._ny*self._nz) self.solution = minimize(self.bound_constraint_u, x0, method='Newton-CG', jac=self.bound_jac_u, hessp=self.bound_hessp_u) def calc_res(self): self.residuals = dict() self.stds = dict() self.residuals['obs'] = np.linalg.norm(self.kernel_op.gtoep.matvec(self.solution)-self.obs_data)**2 self.stds['obs'] = np.std(self.kernel_op.gtoep.matvec(self.solution)-self.obs_data) for key in self.dxyz_constraint.keys(): try: tmp2 = self.solution.reshape(self.dxyz_constraint[key].shape) if ('refer' in self.constraints_val.keys()) and (self.smooth_on == 'm-m0'): tmp2 -= self.constraints_val['refer'].reshape(self.dxyz_constraint[key].shape) tmp2 = np.fft.ifftn(np.fft.fftn(tmp2)*self.dxyz_constraint[key]).real slices = [slice(None)]*tmp2.ndim slices[-1] = slice(0,self.dxyz_spaces[key]) self.residuals[key] = np.linalg.norm(tmp2[tuple(slices)].ravel())**2 self.stds[key] = np.std(tmp2[tuple(slices)].ravel()) except KeyError: pass if 'refer' in self.constraints_val.keys(): self.residuals['refer'] = np.linalg.norm(self.solution.ravel()-self.constraints_val['refer'].ravel())**2 self.stds['refer'] = np.std(self.solution.ravel()-self.constraints_val['refer'].ravel()) @timeit def calc_log_prior_total_det(self): self.log_prior_det_val = 0 self.log_total_det_val = 0 blocks = ['0','1','2','3'] prior_eigs = np.zeros(self._nx*self._ny*self._nz) total_eigs = np.zeros(self._nx*self._ny*self._nz) step = self._nx*self._ny*self._nz//4 try: depth_weight = self._weights['depth'] except KeyError: depth_weight = 1. with h5py.File(self.fname,mode='r') as f: if 'depth' in self._weights.keys(): depth_walsh = f['depth']['0'][:] for i_b,block in enumerate(blocks): tmp_block = np.zeros((step,step)) for dxyz_name in self._smooth_components: try: dxyz_walsh = f[dxyz_name][block][:].reshape(step//self._nz, self._nz, step//self._nz, self._nz) ein_path = np.einsum_path('mi,xiyj,jn->xmyn', depth_walsh.T, dxyz_walsh, depth_walsh, optimize='optimal')[0] tmp_multi = np.einsum('mi,xiyj,jn->xmyn', depth_walsh.T, dxyz_walsh, depth_walsh, optimize=ein_path) tmp_block += depth_weight*self._weights[dxyz_name]*tmp_multi.reshape(step,step) except KeyError: pass if 'refer' in self._weights.keys(): tmp_multi_small = depth_walsh.T@depth_walsh for i in range(step//self._nz): tmp_block[i*self._nz:(i+1)*self._nz, i*self._nz:(i+1)*self._nz] += depth_weight*self._weights['refer']*tmp_multi_small with cp.cuda.Device(2): tmp_block_gpu = cp.asarray(tmp_block,dtype=np.float32) eigs = cp.linalg.eigvalsh(tmp_block_gpu) prior_eigs[i_b*step:(i_b+1)*step] = cp.asnumpy(eigs) self.log_prior_det_val += cp.asnumpy(cp.sum(cp.log(eigs))) tmp_block_gpu = None eigs = None free_gpu() tmp_block += self._weights['obs']*f['kernel'][block][:] with cp.cuda.Device(2): tmp_block_gpu = cp.asarray(tmp_block,dtype=np.float32) eigs = cp.linalg.eigvalsh(tmp_block_gpu) total_eigs[i_b*step:(i_b+1)*step] = cp.asnumpy(eigs) self.log_total_det_val += cp.asnumpy(cp.sum(cp.log(eigs))) tmp_block_gpu = None eigs = None free_gpu() self.log_prior_det_val = cp.asnumpy(self.log_prior_det_val) self.log_total_det_val = cp.asnumpy(self.log_total_det_val) self.eigs = {'prior':prior_eigs,'total':total_eigs} return self.log_prior_det_val,self.log_total_det_val @timeit def calc_log_obs_det(self): self.log_obs_det_val = np.log(self._weights['obs'])*len(self.obs_data) return self.log_obs_det_val @timeit def calc_abic(self): '''-log_prior_det_value+log_total_det-log_obs_det+min_u''' self.calc_log_prior_total_det() self.calc_min_u() self.calc_log_obs_det() self.abic_val = (self.log_total_det_val + self.min_u_val - self.log_prior_det_val - self.log_obs_det_val) return self.abic_val @timeit def para_grad(self,x): pass def u_bound(self): pass def print_summary(self): print('abic values:{}'.format(self.abic_val)) print('log total det:{}'.format(self.log_total_det_val)) print('log prior det:{}'.format(self.log_prior_det_val)) print('log obs det:{}'.format(self.log_obs_det_val)) print('min u:{}'.format(self.min_u_val)) print('std:',end=' ') print(self.stds) print('1/var:',end=' ') print({k:1./v**2 for k,v in self.stds.items()}) print('norms:',end=' ') print(self.residuals)
class FreqInvModel: def __init__(self,conf_file=None,**kwargs): self.confs = {'nzyx':[4,4,4], 'nobsyx':None, 'smooth_components':None, 'depth_scaling':None, 'model_density':None, 'refer_densities':None, 'weights':None, 'source_volume':None, 'obs_area':None, 'data_dir':'./'} confs = dict() if not conf_file is None: with open(conf_file) as f: confs = json.load(f) self.confs = {**self.confs,**confs,**kwargs} self.nz,self.ny,self.nx = self.confs['nzyx'] if self.confs['nobsyx'] is None: self.nobsx = self.nx self.nobsy = self.ny else: self.nobsy,self.nobsx = self.confs['nobsyx'] self.source_volume = self.confs['source_volume'] self.obs_area= self.confs['obs_area'] if self.confs['model_density'] is None: self._model_density = None else: self._model_density = self.confs['model_density'].ravel() self._smooth_components = self.confs['smooth_components'] if self.confs['depth_scaling'] is None: self.depth_scaling = np.ones(self.nz) else: self.depth_scaling = self.confs['depth_scaling'] self.refer_densities = self.confs['refer_densities'] self._weights = self.confs['weights'] self.smop = abic.SmoothOperator() self.kernel_op = None def gen_mesh(self,height = 0): ''' Generate mesh of the model. Args: height (float): height of the observations. ''' shape = (self.nz, self.ny, self.nx) self.mesh = PrismMesh(self.source_volume, shape) self.mesh.addprop('density', self._model_density) # generate obs grid # coordinate: x North-South,y East-West # gridder is in the order: (nx,ny) self.gen_obs_grid(height=height) def gen_obs_grid(self,height=0): if self.obs_area is None: self.obs_area = (self.source_volume[0]+0.5*self.mesh.dims[0], self.source_volume[1]-0.5*self.mesh.dims[0], self.source_volume[2]+0.5*self.mesh.dims[1], self.source_volume[3]-0.5*self.mesh.dims[1]) obs_shape = (self.nobsx, self.nobsy) self.xp, self.yp, self.zp = gridder.regular(self.obs_area, obs_shape, z=height) def _pad_data(self, data, shape): n0 = pftrans._nextpow2(2*shape[0]) n1 = pftrans._nextpow2(2*shape[1]) nx, ny = shape padx = (n0 - nx)//2 pady = (n1 - ny)//2 padded = np.pad(data.reshape(shape), ((padx, padx), (pady, pady)), mode='edge') return padded, padx, pady def gen_kernel(self,gtype='z'): xs = np.array(self.mesh.get_xs()) ys = np.array(self.mesh.get_ys()) zs = np.array(self.mesh.get_zs()) x0 = (xs[:-1] + xs[1:])/2.0 y0 = (ys[:-1] + ys[1:])/2.0 a = np.abs(xs[:-1]-xs[1:])/2.0 b = np.abs(ys[:-1]-ys[1:])/2.0 nx, ny = self.nobsx, self.nobsy xmin, xmax, ymin, ymax = self.obs_area dx = (xmax - xmin)/(nx - 1) dy = (ymax - ymin)/(ny - 1) # Pad the array with the edge values to avoid instability shape = (nx,ny) self.padded, self.padx, self.pady = self._pad_data(self.zp, shape) nxe, nye = self.padded.shape M_left=(nxe-nx)/2+1 M_right=M_left+nx-1 N_down=(nye-ny)/2+1 N_up=N_down+ny-1 XXmin=xmin-dx*(M_left-1) XXmax=xmax+dx*(nxe-M_right) YYmin=ymin-dy*(N_down-1) YYmax=ymax+dy*(nye-N_up) # we store kx and ky as 1d array self.kx = 2*np.pi*np.array(np.fft.fftfreq(self.padded.shape[0], dx)) self.ky = 2*np.pi*np.array(np.fft.fftfreq(self.padded.shape[1], dy)) # kz is 2d array self.kz = np.sqrt(np.add.outer(self.ky**2,self.kx**2)).ravel() self.kxm = np.ma.array(self.kx, mask= self.kx==0) self.kym = np.ma.array(self.ky, mask= self.ky==0) self.kzm = np.ma.array(self.kz, mask= self.kz==0) complex1 = 0+1j ## zs should be depth or coordinate? self.C = -8*np.pi*G*SI2MGAL self.W = np.exp(self.kz*self.zp[0]) if gtype == 'zz': self.W = self.kz * self.W self.dW =((np.exp(-np.outer(self.kz,zs[1:])) -np.exp(-np.outer(self.kz,zs[:-1])))/self.kzm.reshape(-1,1)) self.dW[self.kzm.mask,:] = 0. self.dW[self.kzm.mask,:] += zs[:-1] - zs[1:] self.dW = self.dW.data self.WX = np.exp(complex1*(XXmin-XXmax+xmax+xmin)*self.kx/2.)/dx #self.WX = np.ones_like(self.kx)/dx if gtype == 'zx': self.WX = complex1 * self.kx self.FX = np.sin(np.outer(self.kx,a))/self.kxm.reshape(-1,1) self.FX[self.kxm.mask,:] = 0. self.FX[self.kxm.mask,:] += a self.FX = self.FX * np.exp(-complex1*np.outer(self.kx,x0)) * self.WX.reshape(-1,1) self.FX = self.FX.data self.WY = np.exp(complex1*(YYmin-YYmax+ymax+ymin)*self.ky/2.)/dy # self.WY = np.ones_like(self.ky)/dy if gtype == 'zy': self.WY = complex1 * self.ky self.FY = np.sin(np.outer(self.ky,b))/self.kym.reshape(-1,1) self.FY[self.kym.mask,:] = 0. self.FY[self.kym.mask,:] += b self.FY = self.FY * np.exp(-complex1*np.outer(self.ky,y0)) * self.WY.reshape(-1,1) self.FY = self.FY.data self.solve_prepared = False def _gen_rhs(self): # default pad 0s self.rhs = self.C*self.W*self.obs_freq.T.ravel() self.rhs = self.dW.T*self.rhs # self.rhs = kron_matvec([np.eye(self.nz), # np.conj(self.FY.T), # np.conj(self.FX.T)], # self.rhs.ravel()) self.rhs = kron_matvec([np.conj(self.FY.T), np.conj(self.FX.T)], self.rhs.reshape(self.nz,-1)).ravel() def forward(self,v=None,update=False): ''' Forward modelling gravity and its frequency representation Args: v (ndarray): model density, if not given, use self.model_density update (bool): If False, self.freq and self.obs_field won't be touched. If True, self.freq and self.obs_field will be updated. Returns: obs_field (ndarray): gravity. freq (ndarray): gravity in frequency domain. ''' if v is None: v = self._model_density freq = kron_matvec([self.FY,self.FX],v.reshape(self.nz,-1)) freq = freq * self.dW.T freq = self.W * np.sum(freq,axis=0) freq = self.C * freq.reshape(self.padded.shape[1],self.padded.shape[0]).T obs_field = np.real(np.fft.ifft2(freq)) obs_field = obs_field[self.padx: self.padx + self.nobsx, self.pady: self.pady + self.nobsy].ravel() if update: self.freq = freq self.obs_field = obs_field return freq,obs_field def _prepare_solve(self): if self.solve_prepared: return tmpX = self.FX @ np.conj(self.FX.T) tmpY = [email protected](self.FY.T) tmp = np.kron(tmpY,tmpX) self.woodSUB = np.zeros_like(tmp) for iz in range(self.nz): self.woodSUB += self.dW[:,iz].reshape(-1,1)*tmp*np.conj(self.dW[:,iz].reshape(1,-1)) self.solve_prepared = True def do_linear_solve_quiet(self): self._gen_rhs() if not self.solve_prepared: self._prepare_solve() weight = np.sum(self._weights['refers']) invSUB = self.C**2*weight*self.woodSUB invSUB_diag = np.einsum('ii->i',invSUB) invSUB_diag += weight**2 * 1./self.W**2 invSUB = np.linalg.inv(invSUB) self.solution = kron_matvec([self.FY,self.FX],self.rhs.reshape(self.nz,-1)) self.solution = np.sum(self.solution * self.dW.T, axis=0) self.solution = invSUB @ self.solution self.solution = self.dW.T*self.solution.reshape(1,-1) self.solution = kron_matvec([np.conj(self.FY.T),np.conj(self.FX.T)], self.solution.reshape(self.nz,-1)).ravel() self.solution = self.rhs/weight - self.C**2*self.solution def calc_res(self): pass