cmap=args.cmap)
if epoch % args.log_freq == 0:
logger['psnr_test'].append(psnr)
logger['rmse_test'].append(rmse)
print("Epoch {}: test PSNR: {:.6f}, RMSE: {:.6f}".format(
epoch, psnr, rmse))
tic2 = time.time()
print("Finished training {} epochs using {} seconds".format(
args.epochs, tic2 - tic))
x_axis = np.arange(args.log_freq, args.epochs + args.log_freq,
args.log_freq)
# plot the rmse, r2-score curve and save them in txt
save_stats(args.train_dir, logger, x_axis, 'psnr_train', 'psnr_test',
'rmse_train', 'rmse_test')
args.training_time = tic2 - tic
args.n_params, args.n_layers = module_size(model)
with open(args.run_dir + "/args.txt", 'w') as args_file:
json.dump(vars(args), args_file, indent=4)
except KeyboardInterrupt:
print('Keyboard Interrupt captured...Saving models & training logs')
tic2 = time.time()
torch.save(model.state_dict(),
args.ckpt_dir + "/model_epoch{}.pth".format(epoch))
x_axis = np.arange(args.log_freq, args.epochs + args.log_freq,
args.log_freq)
# plot the rmse, r2-score curve and save them in txt
save_stats(args.train_dir, logger, x_axis, 'psnr_train', 'psnr_test',
print(
f'Epoch {epoch}: training loss: {loss_train:.6f} ' \
f'relative-ux: {relative_ux: .5f}, relative-uy: {relative_uy: .5f}')
logf = open(args.train_dir + "/running_log.txt", 'a')
logf.write(
f'Epoch {epoch}: training loss: {loss_train:.6f}, ' \
f'relative-ux: {relative_ux: .5f}, relative-uy: {relative_uy: .5f} \n')
logf.close()
if epoch % args.log_freq == 0:
logger['loss_train'].append(loss_train)
logger['ux_l2loss'].append(relative_ux)
logger['uy_l2loss'].append(relative_uy)
if epoch % args.ckpt_freq == 0:
torch.save(model, args.ckpt_dir + "/model_epoch{}.pth".format(epoch))
sampledir = args.ckpt_dir + "/model{}".format(epoch)
mkdirs((sampledir))
tic2 = time.time()
print(f'Finished training {args.epochs} epochs with {args.ntrain} data ' \
f'using {(tic2 - tic) / 60:.2f} mins')
metrics = ['loss_train', 'ux_l2loss', 'uy_l2loss']
save_stats(args.train_dir, logger, *metrics)
args.training_time = tic2 - tic
args.n_params, args.n_layers = model.model_size
with open(args.run_dir + "/args.txt", 'w') as args_file:
json.dump(vars(args), args_file, indent=4)
logger['r2_test'].append(r2_test)
logger['mnlp_test'].append(mnlp_test)
print("epoch {}, testing r2: {:.6f}, test mnlp: {}".format(
epoch, r2_test, mnlp_test))
if __name__ == '__main__':
print(
'Start training.........................................................'
)
tic = time()
for epoch in range(1, args.epochs + 1):
svgd.train(epoch, logger)
with torch.no_grad():
test(epoch, logger)
training_time = time() - tic
print(
'Finished training:\n{} epochs\n{} data\n{} samples (SVGD)\n{} seconds'
.format(args.epochs, args.ntrain, args.n_samples, training_time))
# save training results
x_axis = np.arange(args.log_freq, args.epochs + args.log_freq,
args.log_freq)
# plot the rmse, r2-score curve and save them in txt
save_stats(args.train_dir, logger, x_axis)
args.training_time = training_time
args.n_params, args.n_layers = dense_ed._num_parameters_convlayers()
with open(args.run_dir + "/args.txt", 'w') as args_file:
json.dump(vars(args), args_file, indent=4)
def main():
parser = argparse.ArgumentParser(description='CNN to solve PDE')
parser.add_argument('--exp-dir',
type=str,
default='./experiments/solver',
help='color map')
parser.add_argument('--nonlinear',
action='store_true',
default=False,
help='set True for nonlinear PDE')
# data
parser.add_argument('--data-dir',
type=str,
default="./datasets",
help='directory to dataset')
parser.add_argument('--data',
type=str,
default='grf',
choices=['grf', 'channelized', 'warped_grf'],
help='data type')
parser.add_argument('--kle', type=int, default=512, help='# kle terms')
parser.add_argument('--imsize', type=int, default=64, help='image size')
parser.add_argument('--idx',
type=int,
default=8,
help='idx of input, please use 0 ~ 999')
parser.add_argument('--alpha1',
type=float,
default=1.0,
help='coefficient for the squared term')
parser.add_argument('--alpha2',
type=float,
default=1.0,
help='coefficient for the cubic term')
# latent size: (nz, sz, sz)
parser.add_argument('--nz',
type=int,
default=1,
help='# feature maps of latent z')
# parser.add_argument('--sz', type=int, default=16, help='feature map size of latent z')
parser.add_argument('--blocks',
type=list,
default=[8, 6],
help='# layers in each dense block of the decoder')
parser.add_argument('--weight-bound',
type=float,
default=10,
help='weight for boundary condition loss')
parser.add_argument('--lr', type=float, default=0.5, help='learning rate')
parser.add_argument('--epochs',
type=int,
default=500,
help='# epochs to train')
parser.add_argument('--test-freq',
type=int,
default=50,
help='every # epoch to test')
parser.add_argument('--ckpt-freq',
type=int,
default=250,
help='every # epoch to save model')
parser.add_argument('--cmap', type=str, default='jet', help='color map')
parser.add_argument(
'--same-scale',
action='store_true',
help='true for setting noise to be same scale as output')
parser.add_argument('--animate',
action='store_true',
help='true to plot animate figures')
parser.add_argument('--cuda', type=int, default=1, help='cuda number')
parser.add_argument('-v',
'--verbose',
action='store_true',
help='True for versbose output')
args = parser.parse_args()
pprint(vars(args))
device = torch.device(
f"cuda:{args.cuda}" if torch.cuda.is_available() else "cpu")
dataset = f'{args.data}_kle{args.kle}' if args.data == 'grf' else args.data
hyparams = f'{dataset}_idx{args.idx}_dz{args.nz}_blocks{args.blocks}_'\
f'lr{args.lr}_wb{args.weight_bound}_epochs{args.epochs}'
if args.nonlinear:
from utils.fenics import solve_nonlinear_poisson
exp_name = 'conv_mixed_residual_nonlinear'
from models.darcy import conv_constitutive_constraint_nonlinear as constitutive_constraint
hyparams = hyparams + f'_alpha1_{args.alpha1}_alpha2_{args.alpha2}'
else:
exp_name = 'conv_mixed_residual'
from models.darcy import conv_constitutive_constraint as constitutive_constraint
run_dir = args.exp_dir + '/' + exp_name + '/' + hyparams
mkdirs(run_dir)
# load data
assert args.idx < 1000
if args.data == 'grf':
assert args.kle in [512, 128, 1024, 2048]
ntest = 1000 if args.kle == 512 else 1024
hdf5_file = args.data_dir + f'/{args.imsize}x{args.imsize}/kle{args.kle}_lhs{ntest}_test.hdf5'
elif args.data == 'warped_grf':
hdf5_file = args.data_dir + f'/{args.imsize}x{args.imsize}/warped_gp_ng64_n1000.hdf5'
elif args.data == 'channelized':
hdf5_file = args.data_dir + f'/{args.imsize}x{args.imsize}/channel_ng64_n512_test.hdf5'
else:
raise ValueError('No dataset are found for the speficied parameters')
print(f'dataset: {hdf5_file}')
with h5py.File(hdf5_file, 'r') as f:
input_data = f['input'][()]
output_data = f['output'][()]
print(f'input: {input_data.shape}')
print(f'output: {output_data.shape}')
# permeability, (1, 1, 64, 64)
perm_arr = input_data[[args.idx]]
# pressure, flux_hor, flux_ver, (3, 64, 64)
if args.nonlinear:
# solve nonlinear Darcy for perm_arr with FEniCS
output_file = run_dir + '/output_fenics.npy'
if os.path.isfile(output_file):
output_arr = np.load(output_file)
print('Loaded solved output field')
else:
print('Solve nonlinear poisson with FEniCS...')
output_arr = solve_nonlinear_poisson(perm_arr[0, 0], args.alpha1,
args.alpha2, run_dir)
np.save(output_file, output_arr)
else:
output_arr = output_data[args.idx]
print('output shape: ', output_arr.shape)
# model
model = Decoder(args.nz, out_channels=3, blocks=args.blocks).to(device)
print(f'model size: {model.model_size}')
fixed_latent = torch.randn(1, args.nz, 16, 16).to(device) * 0.5
perm_tensor = torch.FloatTensor(perm_arr).to(device)
sobel_filter = SobelFilter(args.imsize, correct=True, device=device)
optimizer = optim.LBFGS(model.parameters(),
lr=args.lr,
max_iter=20,
history_size=50)
logger = {}
logger['loss'] = []
def train(epoch):
model.train()
def closure():
optimizer.zero_grad()
output = model(fixed_latent)
if args.nonlinear:
energy = constitutive_constraint(perm_tensor, output,
sobel_filter, args.alpha1, args.alpha2) \
+ continuity_constraint(output, sobel_filter)
else:
energy = constitutive_constraint(
perm_tensor, output, sobel_filter) + continuity_constraint(
output, sobel_filter)
loss_dirichlet, loss_neumann = boundary_condition(output)
loss_boundary = loss_dirichlet + loss_neumann
loss = energy + loss_boundary * args.weight_bound
loss.backward()
if args.verbose:
print(f'epoch {epoch}: loss {loss.item():6f}, '\
f'energy {energy.item():.6f}, diri {loss_dirichlet.item():.6f}, '\
f'neum {loss_neumann.item():.6f}')
return loss
loss = optimizer.step(closure)
loss_value = loss.item() if not isinstance(loss, float) else loss
logger['loss'].append(loss_value)
print(f'epoch {epoch}: loss {loss_value:.6f}')
if epoch % args.ckpt_freq == 0:
torch.save(model.state_dict(),
run_dir + "/model_epoch{}.pth".format(epoch))
def test(epoch):
if epoch % args.epochs == 0 or epoch % args.test_freq == 0:
output = model(fixed_latent)
output = to_numpy(output)
if args.animate:
i_plot = epoch // args.test_freq
plot_prediction_det_animate2(run_dir,
output_arr,
output[0],
epoch,
args.idx,
i_plot,
plot_fn='imshow',
cmap=args.cmap,
same_scale=args.same_scale)
else:
plot_prediction_det(run_dir,
output_arr,
output[0],
epoch,
args.idx,
plot_fn='imshow',
cmap=args.cmap,
same_scale=args.same_scale)
np.save(run_dir + f'/epoch{epoch}.npy', output[0])
print('start training...')
dryrun = False
tic = time.time()
for epoch in range(1, args.epochs + 1):
if not dryrun:
train(epoch)
test(epoch)
print(
f'Finished optimization for {args.epochs} epochs using {(time.time()-tic)/60:.3f} minutes'
)
save_stats(run_dir, logger, 'loss')
# save input
plt.imshow(perm_arr[0, 0])
plt.colorbar()
plt.savefig(run_dir + '/input.png')
plt.close()
def main():
parser = argparse.ArgumentParser(description='CNN to solve PDE')
parser.add_argument('--exp-dir', type=str, default='./experiments/solver', help='color map')
# data
parser.add_argument('--data-dir', type=str, default="./datasets", help='directory to dataset')
parser.add_argument('--data', type=str, default='grf', choices=['grf', 'channelized', 'warped_grf'], help='data type')
parser.add_argument('--kle', type=int, default=512, help='# kle terms')
parser.add_argument('--imsize', type=int, default=64, help='image size')
parser.add_argument('--idx', type=int, default=8, help='idx of input, please use 0 ~ 999')
parser.add_argument('--alpha1', type=float, default=1.0, help='coefficient for the squared term')
parser.add_argument('--alpha2', type=float, default=1.0, help='coefficient for the cubic term')
parser.add_argument('--dim-hidden', type=int, default=512, help='# nodes in each hidden layer')
parser.add_argument('--layers-hidden', type=int, default=8, help='# hidden layers')
parser.add_argument('--off-grid', action='store_true', help='set True to use colloc ')
parser.add_argument('--n-colloc', type=int, default=4096, help='# collocation points')
parser.add_argument('--weight-bound', type=float, default=10, help='weight for boundary condition loss')
parser.add_argument('--lr', type=float, default=0.5, help='learning rate')
parser.add_argument('--epochs', type=int, default=2000, help='# epochs to train')
parser.add_argument('--test-freq', type=int, default=50, help='every # epoch to test')
parser.add_argument('--ckpt-freq', type=int, default=250, help='every # epoch to save model')
parser.add_argument('--cmap', type=str, default='jet', help='color map')
parser.add_argument('--same-scale', action='store_true', help='true for setting noise to be same scale as output')
parser.add_argument('--animate', action='store_true', help='true to plot animate figures')
parser.add_argument('--cuda', type=int, default=2, help='cuda number')
parser.add_argument('-v', '--verbose', action='store_true', help='True for versbose output')
args = parser.parse_args()
pprint(vars(args))
device = torch.device(f"cuda:{args.cuda}" if torch.cuda.is_available() else "cpu")
exp_name = 'fc_mixed_residual'
dataset = f'{args.data}_kle{args.kle}' if args.data == 'grf' else args.data
hyparams = f'{dataset}_idx{args.idx}_dhid{args.dim_hidden}_lhid{args.layers_hidden}_alpha1_{args.alpha1}_alpha2_{args.alpha2}_'\
f'lr{args.lr}_wb{args.weight_bound}_epochs{args.epochs}_ongrid_{not args.off_grid}_ncolloc{args.n_colloc}'
run_dir = args.exp_dir + '/' + exp_name + '/' + hyparams
mkdirs(run_dir)
# load data
assert args.idx < 1000
if args.data == 'grf':
assert args.kle in [512, 128, 1024, 2048]
ntest = 1000 if args.kle == 512 else 1024
hdf5_file = args.data_dir + f'/{args.imsize}x{args.imsize}/kle{args.kle}_lhs{ntest}_test.hdf5'
elif args.data == 'warped_grf':
hdf5_file = args.data_dir + f'/{args.imsize}x{args.imsize}/warped_gp_ng64_n1000.hdf5'
elif args.data == 'channelized':
assert args.idx < 512
hdf5_file = args.data_dir + f'/{args.imsize}x{args.imsize}/channel_ng64_n512_test.hdf5'
else:
raise ValueError('No dataset are found for the speficied parameters')
print(f'dataset: {hdf5_file}')
with h5py.File(hdf5_file, 'r') as f:
input_data = f['input'][()]
output_data = f['output'][()]
print(f'input: {input_data.shape}')
print(f'output: {output_data.shape}')
# permeability, (1, 1, 64, 64)
perm_arr = input_data[[args.idx]]
# pressure, flux_hor, flux_ver, (3, 64, 64)
output_arr = output_data[args.idx]
def to_tensor_gpu(*numpy_seq):
# x: numpy array --> tensor on GPU
return (torch.FloatTensor(x).to(device) for x in numpy_seq)
# define networks
net_u = CPPN(dim_in=2, dim_out=3, dim_hidden=args.dim_hidden,
layers_hidden=args.layers_hidden).to(device)
print(net_u)
print(net_u._model_size())
optimizer = optim.LBFGS(net_u.parameters(),
lr=args.lr, max_iter=20, history_size=50)
logger = {}
logger['loss'] = []
ngrids = [args.imsize, args.imsize]
sampler = SampleSpatial2d(int(ngrids[0]), int(ngrids[1]))
colloc_on_grid = not args.off_grid
# for batch optimization
x_colloc = sampler.colloc(colloc_on_grid, n_samples=args.n_colloc).to(device)
x_dirichlet = torch.cat((sampler.left(on_grid=False, n_samples=256),
sampler.right(on_grid=False, n_samples=256)), 0).to(device)
y_dirichlet = torch.cat((torch.ones(256, 1), torch.zeros(256, 1)), 0).to(device)
x_neumann = torch.cat((sampler.top(colloc_on_grid),
sampler.bottom(colloc_on_grid)), 0).to(device)
print(sampler.coordinates_no_boundary.shape)
K_true_tensor, = to_tensor_gpu(perm_arr.reshape(-1, 1))
if args.verbose:
print('x_colloc: {}'.format(x_colloc.shape))
print('x_dirc: {}'.format(x_dirichlet.shape))
print('y_dirc: {}'.format(y_dirichlet.shape))
print('x_neumann: {}'.format(x_neumann.shape))
def train(epoch):
net_u.train()
def closure():
optimizer.zero_grad()
loss_colloc = mixed_residual_fc(net_u, x_colloc, K_true_tensor,
args.verbose, rand_colloc=args.off_grid)
# loss_colloc = 0
loss_dirichlet = F.mse_loss(net_u(x_dirichlet)[:, [0]], y_dirichlet)
loss_neumann = neumann_boundary_mixed(net_u, x_neumann)
loss = loss_colloc + args.weight_bound * (loss_dirichlet + loss_neumann)
loss.backward()
if args.verbose:
print(f'epoch {epoch}: colloc {loss_colloc:.6f}, '
f'diri {loss_dirichlet:.6f}, neum {loss_neumann:.6f}')
return loss
loss = optimizer.step(closure)
loss_value = loss.item() if not isinstance(loss, float) else loss
logger['loss'].append(loss_value)
print('epoch {}: loss {:.10f}'.format(epoch, loss_value))
if epoch % args.ckpt_freq == 0:
torch.save(net_u.state_dict(), run_dir + "/model_epoch{}.pth".format(epoch))
def test(epoch):
if epoch % args.epochs == 0 or epoch % args.test_freq == 0:
# plot the solution
xx, yy = np.meshgrid(np.arange(ngrids[0]), np.arange(ngrids[1]))
x_test = xx.flatten()[:, None] / ngrids[1]
y_test = yy.flatten()[:, None] / ngrids[0]
x_test, y_test = to_tensor_gpu(x_test, y_test)
net_u.eval()
x_test.requires_grad = True
y_test.requires_grad = True
xy_test = torch.cat((y_test, x_test), 1)
y_pred = net_u(xy_test)
target = output_arr
# three output of net_u from 0-3 channel: u, flux_y, flux_x
u_pred = y_pred[:, 0].detach().cpu().numpy().reshape(*ngrids)
u_y = y_pred[:, 1].detach().cpu().numpy().reshape(*ngrids)
u_x = y_pred[:, 2].detach().cpu().numpy().reshape(*ngrids)
prediction = np.stack((u_pred, u_x, u_y))
# prediction = y_pred.view(*ngrids, -1).transpose(0, 1).permute(2, 1, 0).detach().cpu().numpy()
if args.animate:
i_plot = epoch // args.test_freq
plot_prediction_det_animate2(run_dir, target, prediction, epoch, args.idx, i_plot,
plot_fn='imshow', cmap=args.cmap, same_scale=args.same_scale)
else:
plot_prediction_det(run_dir, target, prediction, epoch, args.idx,
plot_fn='imshow', cmap=args.cmap, same_scale=args.same_scale)
np.save(run_dir + f'/epoch{epoch}.npy', prediction)
print('start training...')
dryrun = False
tic = time.time()
for epoch in range(1, args.epochs + 1):
if not dryrun:
train(epoch)
test(epoch)
print(f'Finished training {args.epochs} epochs in {(time.time()-tic)/60:.3f} minutes')
save_stats(run_dir, logger, 'loss')
# save input
plt.close()
plt.imshow(np.log(perm_arr[0, 0]))
plt.colorbar()
plt.savefig(run_dir + '/input_logK.png')
plt.close()
# super-resultion one
ngrids = (640, 640)
xx, yy = np.meshgrid(np.arange(ngrids[0]), np.arange(ngrids[1]))
x_test = torch.FloatTensor(np.stack((yy.flatten() / (ngrids[1]-1),
xx.flatten() / (ngrids[0]-1)), 1)).to(device)
net_u.eval()
u_pred = net_u(x_test)
u_pred = u_pred[:, 0].reshape(*ngrids).detach().cpu().numpy()
plt.contourf(u_pred, 65)
plt.colorbar()
plt.savefig(run_dir + '/solution_HR.png')
plt.close()
