import sys, os, time
import numpy as np
import torch
import conf, setenv
configfile = sys.argv[1]
configfile = 'checkpoint/' + configfile + '/options.yaml'
print('configfile: ', configfile)
options = conf.setoptions(configfile=configfile, isload=True)
if not torch.cuda.is_available():
options['--device'] = 'cpu'
if len(sys.argv) > 2:
batch_num = int(sys.argv[2])
else:
batch_num = 0
if len(sys.argv) > 3:
options['--device'] = (sys.argv[3] if sys.argv[3].startswith('cuda')
or sys.argv[3].startswith('cpu') else 'cuda:' +
sys.argv[3])
globalnames, callback, model, data_model, sampling, addnoise = setenv.setenv(
options)
globals().update(globalnames)
outputfile = callback.savepath + '/test-progress'
DATA = torch.load('checkpoint/data-' + dataname + '-' + str(batch_num),
map_location=model.device)
u0_obs, uts_obs, uvars, updatevars = DATA['u0_obs'], DATA['uts_obs'], DATA[
'uvars'], DATA['updatevars']
import sys
import time
import numpy as np
import torch
import conf, setenv
terms0 = []
terms1 = []
coeffs0 = np.zeros((40))
coeffs1 = np.zeros((40))
options = {}
options['--name'] = sys.argv[1]
configfile = 'checkpoint/' + options['--name'] + '/options.yaml'
options = conf.setoptions(argv=None,
kw=None,
configfile=configfile,
isload=True)
options['--device'] = 'cpu'
block = sys.argv[2]
globalnames, callback, model, data_model, sampling, addnoise = setenv.setenv(
options)
callback.load(int(block))
tsym, csym = model.poly0.coeffs(calprec=8)
terms0 = tsym
coeffs0[:min(len(tsym), 40)] = csym[:min(len(tsym), 40)]
tsym, csym = model.poly1.coeffs(calprec=8)
terms1 = tsym
coeffs1[:min(len(tsym), 40)] = csym[:min(len(tsym), 40)]
torch.save(
def setenv(options):
"""
set training & testing environment
Returns:
globalnames(dict): variables need to be exported to training & testing script
callback(function class): callback function for optimizer
model(torch.nn.Module): PDE-Net, a torch forward neural network
data_model(torch.nn.Module): a torch module for data generation
sampling,addnoise(callable function): data down sample and add noise to data
"""
globalnames, callback, model = _set_model(options)
if options['--dataname'] == 'None':
dataoptions = conf.setoptions(configfile='checkpoint/'+
options['--dataname']+'/options.yaml',
isload=True)
dataoptions['--start_from'] = 80
assert options['--cell_num']%dataoptions['--cell_num'] == 0
dataoptions['--device'] = options['--device']
dataoptions['--dtype'] = options['--dtype']
_,_,data_model = _set_model(dataoptions)
data_model.tiling = options['--cell_num']//dataoptions['--cell_num']
mesh_size = list(m*globalnames['zoom'] for m in globalnames['mesh_size'])
mesh_bound = globalnames['mesh_bound']
viscosity = globalnames['viscosity']
dx = globalnames['cell_num']*globalnames['eps']/mesh_size[0]
if options['--dataname'].upper() == 'BURGERS':
max_dt = globalnames['max_dt']
data_model = burgers2d.BurgersTime2d(max_dt=max_dt,
mesh_size=mesh_size,
mesh_bound=mesh_bound,
viscosity=viscosity,
timescheme=globalnames['data_timescheme'],
)
elif options['--dataname'].upper() == 'HEAT':
max_dt = globalnames['max_dt']
data_model = cde2d.Heat(max_dt=max_dt,
mesh_size=mesh_size,
mesh_bound=mesh_bound,
timescheme=globalnames['data_timescheme']
)
data_model.coe[0,2] = data_model.coe[2,0] = viscosity
elif options['--dataname'].upper() == 'REACTIONDIFFUSION':
max_dt = globalnames['max_dt']
data_model = rd2d.RDTime2d(max_dt=max_dt,
mesh_size=mesh_size,
mesh_bound=mesh_bound,
viscosity=viscosity,
beta=1,
timescheme=globalnames['data_timescheme']
)
elif options['--dataname'].upper() == 'CDR':
max_dt = globalnames['max_dt']
data_model = cdr2d.CDRTime2d(max_dt=max_dt,
mesh_size=mesh_size,
mesh_bound=mesh_bound,
viscosity=viscosity,
beta=1,
timescheme=globalnames['data_timescheme'],
)
data_model.to(device=model.device)
if globalnames['dtype'] == torch.float64:
data_model.double()
else:
data_model.float()
sampling = transform.Compose(
transform.DownSample(mesh_size=globalnames['mesh_size']),
)
addnoise = transform.AddNoise(start_noise=options['--start_noise'], end_noise=options['--end_noise'])
return globalnames, callback, model, data_model, sampling, addnoise
import conf, setenv, initparameters
import matplotlib.pyplot as plt
import aTEAM.pdetools as pdetools
from pltutils import *
#%%
options_ref = configfile = None
options_ref = conf.default_options()
options_ref['--dataname'] = 'heat'
options_ref['--viscosity'] = 0.1
options_ref['--dt'] = 1e-2
options_ref['--max_dt'] = 1e-2 / 16
options_ref['--zoom'] = 4
options_ref['--batch_size'] = 2
options_ref['--data_timescheme'] = 'rk2'
options_ref['--channel_names'] = 'u'
options_ref = conf.setoptions(argv=None, kw=options_ref, configfile=configfile)
if torch.cuda.is_available():
options_ref['--device'] = 'cuda'
else:
options_ref['--device'] = 'cpu'
globalnames_ref, callback_ref, model_ref, data_model_ref, sampling_ref, addnoise_ref = setenv.setenv(
options_ref)
globals().update(globalnames_ref)
# initialization of parameters
initparameters.initkernels(model_ref)
initparameters.initexpr(model_ref, viscosity=viscosity, pattern=dataname)
# model_ref.polys[k].coeffs(iprint=1)
# # data transform
# '--start_noise':0.001,
# '--end_noise':0.001,
# # others
# '--stablize':0.0,
# '--sparsity':0.005,
# '--momentsparsity':0.001,
# '--npseed':-1,
# '--torchseed':-1,
# '--maxiter':2000,
# '--recordfile':'None',
# '--recordcycle':200,
# '--savecycle':-1,
# '--start_from':-1,
# }
options = conf.setoptions(argv=sys.argv[1:], kw=kw, configfile=None)
print(options)
globalnames, callback, model, data_model, sampling, addnoise = setenv.setenv(
options)
globals().update(globalnames)
torch.cuda.manual_seed_all(torchseed)
torch.manual_seed(torchseed)
np.random.seed(npseed)
# initialization of parameters
if start_from < 0:
initparameters.initkernels(model, scheme=scheme)
# initparameters.renormalize(model, u0)
