def main(data,gf,param,outfile): ''' Nt: number of time steps Nx: number of positions Dx: spatial dimensions of coordinates and displacements Ns: number of slip basis functions per slip direction Ds: number of slip directions Nv: number of fluidity basis functions total: number of state parameters (Ns*Ds + Nv + 2*Nx*Dx) Parameters ---------- data: \ mask : (Nt,Nx) boolean array \ mean : (Nt,Nx,Dx) array \ covariance : (Nt,Nx,Dx,Dx) array \ metadata \ position : (Nx,Dx) array time : (Nt,) array prior: \ mean : (total,) array \ covariance : (total,total) array gf: \ elastic : (Ns,Ds,Nx,Dx) array \ viscoelastic : (Ns,Ds,Dv,Nx,Dx) array \ metadata \ position : (Nx,Dx) array reg: \ regularization : (*,total) array params: user parameter dictionary Returns ------- out: \ slip_integral \ mean : (Nt,Ns,Ds) array \ uncertainty :(Nt,Ns,Ds) array \ slip \ mean : (Nt,Ns,Ds) array \ uncertainty : (Nt,Ns,Ds) array \ slip_derivative \ mean : (Nt,Ns,Ds) array \ uncertainty : (Nt,Ns,Ds) array \ fluidity \ mean : (Nv,) array \ uncertainty : (Nv,) array \ secular_velocity \ mean : (Nx,Dx) array \ uncertainty : (Nx,Dx) array \ baseline_displacement \ mean : (Nx,Dx) array \ uncertainty : (Nx,Dx) array ''' F = gf['slip'][...] G = gf['fluidity'][...] coseismic_times = np.array(param['coseismic_times']) afterslip_start_times = param['afterslip_start_times'] afterslip_end_times = param['afterslip_end_times'] afterslip_times = zip(afterslip_start_times,afterslip_end_times) afterslip_times = np.array(afterslip_times) time = data['time'][:] slip_scale = 1.0 # meter relax_scale = 1.0 # year time_scale = np.std(time) time_shift = np.mean(time) # years disp_scale = 1.0 # meter # shift is different for each time series disp_shift = np.mean(data['mean'],0) # scale greens functions # F is originally in meters disp per meter slip F /= disp_scale/slip_scale # G is originally in meters/year disp per meter slip*fluidity G /= (disp_scale/time_scale)/(slip_scale/relax_scale) time -= time_shift time /= time_scale coseismic_times -= time_shift coseismic_times /= time_scale afterslip_times -= time_shift afterslip_times /= time_scale param['initial_slip_variance'] /= slip_scale**2 param['fluidity_variance'] *= relax_scale**2 param['secular_velocity_variance'] /= (disp_scale/time_scale)**2 param['baseline_displacement_variance'] /= disp_scale**2 # define slip functions and slip jacobian here slip_func,slip_jac = steps_and_ramps(coseismic_times, afterslip_times) Nst = len(coseismic_times) + len(afterslip_start_times) Ns,Ds,Nv,Nx,Dx = np.shape(G) Nt = len(time) # check for consistency between input assert data['mean'].shape == (Nt,Nx,Dx) assert data['variance'].shape == (Nt,Nx,Dx) assert F.shape == (Ns,Ds,Nx,Dx) assert G.shape == (Ns,Ds,Nv,Nx,Dx) p = state_parser(Nst,Ns,Ds,Nv,Nx,Dx) if param['solver'] == 'bvls': solver = modest.bvls upper_bound = 1e6*np.ones(p['total']) lower_bound = -1e6*np.ones(p['total']) # all inferred fluidities will be positive lower_bound[p['fluidity']] = 0 # all inferred left lateral slip will be positive left_lateral_indices = np.array(p['slip'][:,:,0],copy=True) thrust_indices = np.array(p['slip'][:,:,1],copy=True) upper_bound[left_lateral_indices] = 0.0 solver_args = (lower_bound,upper_bound) solver_kwargs = {} if param['solver'] == 'lstsq': solver = modest.lstsq solver_args = () solver_kwargs = {} elif param['solver'] == 'lsmr': solver = modest.lsmr solver_args = () solver_kwargs = {} elif param['solver'] == 'lgmres': solver = modest.lgmres solver_args = () solver_kwargs = {'tol':1e-8,'maxiter':1000} elif param['solver'] == 'dgs': solver = modest.dgs solver_args = () solver_kwargs = {} #fprior = prior.create_formatted_prior(param,p,slip_model='parameterized') Xprior,Cprior = prior.create_prior(param,p,slip_model='parameterized') reg_matrix = reg.create_regularization(param,p,slip_model='parameterized') reg_rows = len(reg_matrix) setup_output_file(outfile,p, data['name'][...], data['position'][...], data['time'][...]) Xprior,Cprior = modest.nonlin_lstsq( regularization, np.zeros(reg_rows), Xprior, data_covariance=np.eye(reg_rows), prior_covariance=Cprior, system_args=(reg_matrix,), jacobian=regularization_jacobian, jacobian_args=(reg_matrix,), maxitr=param['maxitr'], solver=solver, solver_args=solver_args, solver_kwargs=solver_kwargs, LM_damping=True, output=['solution','solution_covariance']) time_indices = range(Nt) block_time_indices = modest.misc.divide_list(time_indices,param['time_blocks']) for i in block_time_indices: outfile['data/mean'][i,...] = data['mean'][i,...] outfile['data/mask'][i,...] = data['mask'][i,...] outfile['data/covariance'][i,...] = data['variance'][i,...] di = data['mean'][i,...] di -= disp_shift di /= disp_scale di_mask = np.array(data['mask'][i,:],dtype=bool) # expand to three dimensions di_mask = np.repeat(di_mask[...,None],3,-1) di = di[~di_mask] Cdi = data['variance'][i,...] Cdi /= disp_scale**2 Cdi = Cdi[~di_mask] Xprior,Cprior = modest.nonlin_lstsq( observation, di, Xprior, data_covariance=Cdi, prior_covariance=Cprior, system_args=(time[i],F,G,p,slip_func,di_mask), jacobian=observation_jacobian, jacobian_args=(time[i],F,G,p,slip_func,slip_jac,di_mask), solver=solver, solver_args=solver_args, solver_kwargs=solver_kwargs, maxitr=param['maxitr'], LM_damping=True, LM_param=1.0, rtol=1e-2, atol=1e-2, output=['solution','solution_covariance']) post_mean_scaled,post_cov_scaled = Xprior,Cprior post_mean = np.copy(post_mean_scaled) post_mean[p['baseline_displacement']] *= disp_scale post_mean[p['baseline_displacement']] += disp_shift post_mean[p['secular_velocity']] *= (disp_scale/time_scale) post_mean[p['slip']] *= slip_scale post_mean[p['fluidity']] /= relax_scale for i in range(Nt): outfile['state/all'][i,:] = post_mean outfile['state/baseline_displacement'][i,...] = post_mean[p['baseline_displacement']] outfile['state/secular_velocity'][i,...] = post_mean[p['secular_velocity']] outfile['state/slip'][i,...] = slip_func(post_mean[p['slip']],time[i]) outfile['state/slip_derivative'][i,...] = slip_func(post_mean[p['slip']],time[i],diff=1) outfile['state/fluidity'][i,...] = post_mean[p['fluidity']] # compute predicted data logger.info('computing predicted data') error = 0.0 count = 0 for i in range(Nt): predicted = observation_t(post_mean_scaled, time[i], F,G,p,slip_func) # change back to meters predicted *= disp_scale predicted += disp_shift residual = outfile['data/mean'][i,...] - predicted covariance = outfile['data/covariance'][i,...] data_mask = np.array(outfile['data/mask'][i,...],dtype=bool) error += L2(residual,covariance,data_mask) count += np.sum(~data_mask) outfile['predicted/mean'][i,...] = predicted mask = np.zeros(p['total']) mask[p['secular_velocity']] = 1.0 mask[p['baseline_displacement']] = 1.0 mask_post_mean = post_mean_scaled*mask tectonic = observation_t(mask_post_mean, time[i], F,G,p, slip_func) tectonic *= disp_scale tectonic += disp_shift outfile['tectonic/mean'][i,...] = tectonic mask = np.zeros(p['total']) mask[p['slip']] = 1.0 mask_post_mean = post_mean_scaled*mask elastic = observation_t(mask_post_mean, time[i], F,G,p, slip_func) elastic *= disp_scale outfile['elastic/mean'][i,...] = elastic visc = (outfile['predicted/mean'][i,...] - outfile['tectonic/mean'][i,...] - outfile['elastic/mean'][i,...]) outfile['viscous/mean'][i,...] = visc logger.info('total RMSE: %s' % np.sqrt(error/count)) return
def main(data,gf,param,outfile): ''' Nt: number of time steps Nx: number of positions Dx: spatial dimensions of coordinates and displacements Ns: number of slip basis functions per slip direction Ds: number of slip directions Nv: number of fluidity basis functions total: number of state parameters (Ns*Ds + Nv + 2*Nx*Dx) Parameters ---------- data: \ mask : (Nt,Nx) boolean array \ mean : (Nt,Nx,Dx) array \ covariance : (Nt,Nx,Dx,Dx) array \ metadata \ position : (Nx,Dx) array time : (Nt,) array prior: \ mean : (total,) array \ covariance : (total,total) array gf: \ elastic : (Ns,Ds,Nx,Dx) array \ viscoelastic : (Ns,Ds,Dv,Nx,Dx) array \ metadata \ position : (Nx,Dx) array reg: \ regularization : (*,total) array params: user parameter dictionary Returns ------- out: \ slip_integral \ mean : (Nt,Ns,Ds) array \ uncertainty :(Nt,Ns,Ds) array \ slip \ mean : (Nt,Ns,Ds) array \ uncertainty : (Nt,Ns,Ds) array \ slip_derivative \ mean : (Nt,Ns,Ds) array \ uncertainty : (Nt,Ns,Ds) array \ fluidity \ mean : (Nv,) array \ uncertainty : (Nv,) array \ secular_velocity \ mean : (Nx,Dx) array \ uncertainty : (Nx,Dx) array \ baseline_displacement \ mean : (Nx,Dx) array \ uncertainty : (Nx,Dx) array ''' # set attributes of outfile outfile.attrs.update(param) outfile.attrs.update(gf.attrs) outfile.attrs.update(data.attrs) F = gf['slip'][...] G = gf['fluidity'][...] alpha = param['slip_acceleration_variance'] jump_factor = param['jump_variance_factor'] jump_times = param['jump_times'] time = data['time'][:] Ns,Ds,Nv,Nx,Dx = np.shape(G) Nt = len(time) # check for consistency between input assert np.shape(data['mean']) == (Nt,Nx,Dx) assert np.shape(data['covariance']) == (Nt,Nx,Dx,Dx) assert np.shape(F) == (Ns,Ds,Nx,Dx) assert np.shape(G) == (Ns,Ds,Nv,Nx,Dx) p = state_parser(Ns,Ds,Nv,Nx,Dx) fprior = prior.create_formatted_prior(param,p) Xprior,Cprior = prior.create_prior(fprior,p) freg = reg.create_formatted_regularization(param,p) reg_matrix = reg.create_regularization(freg,p) reg_rows = len(reg_matrix) kalman = modest.KalmanFilter( Xprior,Cprior, observation_augmented, obs_args=(F,G,p,reg_matrix), ojac=observation_jacobian_augmented, ojac_args=(F,G,p,reg_matrix), trans=transition, trans_args=(p,), tjac=transition_jacobian, tjac_args=(p,), pcov=process_covariance, pcov_args=(alpha,p), core=False, solver_kwargs={'maxitr':param['maxitr']}, temp_file=param['history_file']) # prepare outfile file # copy over data file outfile.create_dataset('data/mean',shape=np.shape(data['mean'])) outfile.create_dataset('data/mask',shape=np.shape(data['mask']),dtype=bool) outfile.create_dataset('data/covariance',shape=np.shape(data['covariance'])) outfile['data/name'] = data['name'][...] outfile['data/position'] = data['position'][...] outfile['data/time'] = data['time'][...] outfile.create_dataset('predicted/mean',shape=np.shape(data['mean'])) outfile.create_dataset('predicted/covariance',shape=np.shape(data['covariance'])) outfile['predicted/name'] = data['name'][...] outfile['predicted/mask'] = np.array(0.0*data['mask'][...],dtype=bool) outfile['predicted/position'] = data['position'][...] outfile['predicted/time'] = data['time'][...] outfile.create_dataset('tectonic/mean',shape=np.shape(data['mean'])) outfile.create_dataset('tectonic/covariance',shape=np.shape(data['covariance'])) outfile['tectonic/name'] = data['name'][...] outfile['tectonic/mask'] = np.array(0.0*data['mask'][...],dtype=bool) outfile['tectonic/position'] = data['position'][...] outfile['tectonic/time'] = data['time'][...] outfile.create_dataset('elastic/mean',shape=np.shape(data['mean'])) outfile.create_dataset('elastic/covariance',shape=np.shape(data['covariance'])) outfile['elastic/name'] = data['name'][...] outfile['elastic/mask'] = np.array(0.0*data['mask'][...],dtype=bool) outfile['elastic/position'] = data['position'][...] outfile['elastic/time'] = data['time'][...] outfile.create_dataset('viscous/mean',shape=np.shape(data['mean'])) outfile.create_dataset('viscous/covariance',shape=np.shape(data['covariance'])) outfile['viscous/name'] = data['name'][...] outfile['viscous/mask'] = np.array(0.0*data['mask'][...],dtype=bool) outfile['viscous/position'] = data['position'][...] outfile['viscous/time'] = data['time'][...] a = np.ones((Nt,Nx)) b = np.eye(3) c = 1e-10*np.einsum('ij,kl',a,b) outfile['predicted/covariance'][...] = c outfile['tectonic/covariance'][...] = c outfile['elastic/covariance'][...] = c outfile['viscous/covariance'][...] = c outfile.create_dataset('state/all',shape=(Nt,p['total'])) for k in ['baseline_displacement', 'secular_velocity', 'slip_integral', 'slip', 'slip_derivative', 'fluidity']: outfile.create_dataset('state/' + k,shape=(Nt,)+np.shape(p[k])) outfile['state/time'] = data['time'][...] logger.info('starting Kalman filter iterations') for i in range(Nt): outfile['data/mean'][i,...] = data['mean'][i,...] outfile['data/mask'][i,...] = data['mask'][i,...] outfile['data/covariance'][i,...] = data['covariance'][i,...] di = flat_data(data['mean'][i,...]) di = np.hstack((di,np.zeros(reg_rows))) di_mask = np.array(data['mask'][i,:],dtype=bool) # expand to three dimensions di_mask = np.repeat(di_mask[:,None],3,1) di_mask = flat_data(di_mask) di_mask = np.hstack((di_mask,np.zeros(reg_rows,dtype=bool))) Cdi = flat_cov(data['covariance'][i,...]) Cdi = scipy.linalg.block_diag(Cdi,np.eye(reg_rows)) if np.all(time[i] < jump_times): kalman.pcov_args = (0.0*alpha,p) kalman.next(di,Cdi,time[i],mask=di_mask) kalman.pcov_args = (alpha,p) elif (i != 0) & np.any((time[i] > jump_times) & (time[i-1] <= jump_times)): logger.info('increasing slip variance by %sx when updating ' 'from t=%s to t=%s' % (jump_factor,time[i-1],time[i])) kalman.pcov_args = (jump_factor*alpha,p) kalman.next(di,Cdi,time[i],mask=di_mask) kalman.pcov_args = (alpha,p) else: kalman.next(di,Cdi,time[i],mask=di_mask) print(kalman.get_posterior()[0][p['fluidity']]) logger.info('starting Rauch-Tung-Striebel smoothing') kalman.smooth() for i in range(Nt): outfile['state/all'][i,:] = kalman.history['smooth'][i,:] for k in ['baseline_displacement', 'secular_velocity', 'slip_integral', 'slip', 'slip_derivative', 'fluidity']: outfile['state/' + k][i,...] = np.array(kalman.history['smooth'])[i,p[k]] # compute predicted data logger.info('computing predicted data') error = 0.0 for i in range(Nt): predicted = observation(outfile['state/all'][i], time[i], F,G,p,flatten=False) residual = outfile['data/mean'][i,...] - predicted covariance = outfile['data/covariance'][i,...] data_mask = np.array(outfile['data/mask'][i,...],dtype=bool) error += RMSE(residual,covariance,data_mask) outfile['predicted/mean'][i,...] = predicted mask = np.zeros(p['total']) mask[p['secular_velocity']] = 1.0 mask[p['baseline_displacement']] = 1.0 state = outfile['state/all'][i]*mask outfile['tectonic/mean'][i,...] = observation( state, time[i], F,G,p, flatten=False) mask = np.zeros(p['total']) mask[p['slip']] = 1.0 state = outfile['state/all'][i]*mask outfile['elastic/mean'][i,...] = observation( state, time[i], F,G,p, flatten=False) visc = (outfile['predicted/mean'][i,...] - outfile['tectonic/mean'][i,...] - outfile['elastic/mean'][i,...]) outfile['viscous/mean'][i,...] = visc logger.info('total RMSE: %s' % error) return